diff --git a/README.md b/README.md
index fe548354..b3d74fa2 100644
--- a/README.md
+++ b/README.md
@@ -1,19 +1,46 @@
+
 # diffusion2D
 
-## Instructions for students
+## Project description
 
-Please follow the instructions in [pypi_exercise.md](https://github.com/Simulation-Software-Engineering/Lecture-Material/blob/main/03_building_and_packaging/pypi_exercise.md).
+This code solves the diffusion equation in 2D over a square domain which is at a certain temperature and a circular disc at the center which is at a higher temperature. This code solves the diffusion equation using the Finite Difference Method. The thermal diffusivity and initial conditions of the system can be changed by the user. The code produces four plots at various timepoints of the simulation. The diffusion process can be clearly observed in these plots.
+Take a few minutes to play around with parameters dx, dy and D in the solver file and observe how the value of dt and the output changes. Do you notice if the code takes more or less time to finish the computation? This tuning is only for you to understand the underlying physical phenomenon and not part of the evaluation.
 
-The code used in this exercise is based on [Chapter 7 of the book "Learning Scientific Programming with Python"](https://scipython.com/book/chapter-7-matplotlib/examples/the-two-dimensional-diffusion-equation/).
+## Installing the package
 
-## Project description
+There are two ways to install this package:
 
-## Installing the package
+### Installing from Github
+
+1. `git clone https://github.com/MarcelWolkober/diffusion2D`
+
+2. `cd diffusion2D`
+
+3. `pip install .`
 
 ### Using pip3 to install from PyPI
 
+`pip3 install --index-url https://test.pypi.org/simple/ wolkobml_diffusion2d --extra-index-url https://pypi.org/simple`
+
 ### Required dependencies
 
+The required dependencies are [Numpy](https://numpy.org) and [Matplotlib](https://matplotlib.org) which are automatically installed.
+
 ## Running this package
 
+Use the provided `solve()` function in python:
+
+```python
+from wolkobml_diffusion2d.diffusion2d import solve
+solve(dx = 0.1, dy = 0.1, D = 4)
+```
+
+It contains three parameter, which can be adjusted:
+
+- `dx` intervals in x-direction
+- `dy` intervals in y-direction
+- `D` thermal diffusivity
+
 ## Citing
+
+This is a student project forked from <https://github.com/Simulation-Software-Engineering/diffusion2D>.
diff --git a/diffusion2d.py b/diffusion2d.py
deleted file mode 100644
index c0c6083a..00000000
--- a/diffusion2d.py
+++ /dev/null
@@ -1,81 +0,0 @@
-"""
-Solving the two-dimensional diffusion equation
-
-Example acquired from https://scipython.com/book/chapter-7-matplotlib/examples/the-two-dimensional-diffusion-equation/
-"""
-
-import numpy as np
-import matplotlib.pyplot as plt
-
-# plate size, mm
-w = h = 10.
-# intervals in x-, y- directions, mm
-dx = dy = 0.1
-# Thermal diffusivity of steel, mm^2/s
-D = 4.
-
-# Initial cold temperature of square domain
-T_cold = 300
-
-# Initial hot temperature of circular disc at the center
-T_hot = 700
-
-# Number of discrete mesh points in X and Y directions
-nx, ny = int(w / dx), int(h / dy)
-
-# Computing a stable time step
-dx2, dy2 = dx * dx, dy * dy
-dt = dx2 * dy2 / (2 * D * (dx2 + dy2))
-
-print("dt = {}".format(dt))
-
-u0 = T_cold * np.ones((nx, ny))
-u = u0.copy()
-
-# Initial conditions - circle of radius r centred at (cx,cy) (mm)
-r = min(h, w) / 4.0
-cx = w / 2.0
-cy = h / 2.0
-r2 = r ** 2
-for i in range(nx):
-    for j in range(ny):
-        p2 = (i * dx - cx) ** 2 + (j * dy - cy) ** 2
-        if p2 < r2:
-            u0[i, j] = T_hot
-
-
-def do_timestep(u_nm1, u, D, dt, dx2, dy2):
-    # Propagate with forward-difference in time, central-difference in space
-    u[1:-1, 1:-1] = u_nm1[1:-1, 1:-1] + D * dt * (
-            (u_nm1[2:, 1:-1] - 2 * u_nm1[1:-1, 1:-1] + u_nm1[:-2, 1:-1]) / dx2
-            + (u_nm1[1:-1, 2:] - 2 * u_nm1[1:-1, 1:-1] + u_nm1[1:-1, :-2]) / dy2)
-
-    u_nm1 = u.copy()
-    return u_nm1, u
-
-
-# Number of timesteps
-nsteps = 101
-# Output 4 figures at these timesteps
-n_output = [0, 10, 50, 100]
-fig_counter = 0
-fig = plt.figure()
-
-# Time loop
-for n in range(nsteps):
-    u0, u = do_timestep(u0, u, D, dt, dx2, dy2)
-
-    # Create figure
-    if n in n_output:
-        fig_counter += 1
-        ax = fig.add_subplot(220 + fig_counter)
-        im = ax.imshow(u.copy(), cmap=plt.get_cmap('hot'), vmin=T_cold, vmax=T_hot)  # image for color bar axes
-        ax.set_axis_off()
-        ax.set_title('{:.1f} ms'.format(n * dt * 1000))
-
-# Plot output figures
-fig.subplots_adjust(right=0.85)
-cbar_ax = fig.add_axes([0.9, 0.15, 0.03, 0.7])
-cbar_ax.set_xlabel('$T$ / K', labelpad=20)
-fig.colorbar(im, cax=cbar_ax)
-plt.show()
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 00000000..cc50cc88
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,23 @@
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "wolkobml_diffusion2d"
+version = "0.0.3"
+description = "Solving the two-dimensional diffusion equation"
+keywords = ["simulation", "diffusion"]
+readme = "README.md"
+requires-python = ">=3.6"
+license = { file = "LICENSE" }
+authors = [
+    { name = "Marcel Wolkober", email = "st163937@stud.uni-stuttgart.de" }
+]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Operating System :: OS Independent",
+]
+dependencies = ["numpy", "matplotlib"]
+
+[project.urls]
+repository = "https://github.com/MarcelWolkober/diffusion2D"
\ No newline at end of file
diff --git a/setup.py b/setup.py
new file mode 100644
index 00000000..4a15e68f
--- /dev/null
+++ b/setup.py
@@ -0,0 +1,4 @@
+import setuptools
+
+if __name__ == "__main__":
+    setuptools.setup()
\ No newline at end of file
diff --git a/wolkobml_diffusion2d/__init__.py b/wolkobml_diffusion2d/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/wolkobml_diffusion2d/diffusion2d.py b/wolkobml_diffusion2d/diffusion2d.py
new file mode 100644
index 00000000..2fd8fa68
--- /dev/null
+++ b/wolkobml_diffusion2d/diffusion2d.py
@@ -0,0 +1,71 @@
+"""
+Solving the two-dimensional diffusion equation
+
+Example acquired from https://scipython.com/book/chapter-7-matplotlib/examples/the-two-dimensional-diffusion-equation/
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+from .output import create_plot, output_plots
+
+def solve(dx = 0.1, dy = 0.1, D = 4):
+    # plate size, mm
+    w = h = 10.
+    
+    # Initial cold temperature of square domain
+    T_cold = 300
+
+    # Initial hot temperature of circular disc at the center
+    T_hot = 700
+
+    # Number of discrete mesh points in X and Y directions
+    nx, ny = int(w / dx), int(h / dy)
+
+    # Computing a stable time step
+    dx2, dy2 = dx * dx, dy * dy
+    dt = dx2 * dy2 / (2 * D * (dx2 + dy2))
+
+    print("dt = {}".format(dt))
+
+    u0 = T_cold * np.ones((nx, ny))
+    u = u0.copy()
+
+    # Initial conditions - circle of radius r centred at (cx,cy) (mm)
+    r = min(h, w) / 4.0
+    cx = w / 2.0
+    cy = h / 2.0
+    r2 = r ** 2
+    for i in range(nx):
+        for j in range(ny):
+            p2 = (i * dx - cx) ** 2 + (j * dy - cy) ** 2
+            if p2 < r2:
+                u0[i, j] = T_hot
+
+
+    def do_timestep(u_nm1, u, D, dt, dx2, dy2):
+        # Propagate with forward-difference in time, central-difference in space
+        u[1:-1, 1:-1] = u_nm1[1:-1, 1:-1] + D * dt * (
+                (u_nm1[2:, 1:-1] - 2 * u_nm1[1:-1, 1:-1] + u_nm1[:-2, 1:-1]) / dx2
+                + (u_nm1[1:-1, 2:] - 2 * u_nm1[1:-1, 1:-1] + u_nm1[1:-1, :-2]) / dy2)
+
+        u_nm1 = u.copy()
+        return u_nm1, u
+
+
+    # Number of timesteps
+    nsteps = 101
+    # Output 4 figures at these timesteps
+    n_output = [0, 10, 50, 100]
+    fig_counter = 0
+    fig = plt.figure()
+
+    # Time loop
+    for n in range(nsteps):
+        u0, u = do_timestep(u0, u, D, dt, dx2, dy2)
+
+        # Create figure
+        if n in n_output:
+            im, fig_counter =  create_plot(fig, u, n, dt, T_cold, T_hot,  fig_counter)
+
+    # Plot output figures
+    output_plots(fig, im)
\ No newline at end of file
diff --git a/wolkobml_diffusion2d/output.py b/wolkobml_diffusion2d/output.py
new file mode 100644
index 00000000..293d7d53
--- /dev/null
+++ b/wolkobml_diffusion2d/output.py
@@ -0,0 +1,16 @@
+import matplotlib.pyplot as plt
+
+def create_plot(fig, u, n, dt, T_cold, T_hot,  fig_counter):
+    fig_counter += 1
+    ax = fig.add_subplot(220 + fig_counter)
+    im = ax.imshow(u.copy(), cmap=plt.get_cmap('hot'), vmin=T_cold, vmax=T_hot)  # image for color bar axes
+    ax.set_axis_off()
+    ax.set_title('{:.1f} ms'.format(n * dt * 1000))
+    return im, fig_counter
+
+def output_plots(fig, im):
+    fig.subplots_adjust(right=0.85)
+    cbar_ax = fig.add_axes([0.9, 0.15, 0.03, 0.7])
+    cbar_ax.set_xlabel('$T$ / K', labelpad=20)
+    fig.colorbar(im, cax=cbar_ax)
+    plt.show()