diff --git a/.gitignore b/.gitignore
index fa358a0..2af3a7d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,3 +5,4 @@
 .DS_Store
 Thumbs.db
 .ipynb_checkpoints/
+.venv/
diff --git a/pyproject.toml b/pyproject.toml
index 1705ffc..5b33d98 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -14,5 +14,8 @@ dependencies = [
     "numpy",
 ]
 
-[tool.setuptools]
-py-modules = ["diffusion_model"]
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+pythonpath = ["src"]
diff --git a/diffusion_model.py b/src/diffusion_model/__init__.py
similarity index 83%
rename from diffusion_model.py
rename to src/diffusion_model/__init__.py
index f23e9ad..20b543c 100644
--- a/diffusion_model.py
+++ b/src/diffusion_model/__init__.py
@@ -1,5 +1,3 @@
-import os
-import sys
 import tomllib
 
 import numpy as np
@@ -95,23 +93,3 @@ def load_params_from_path(filepath):
     with open(filepath, "rb") as stream:
         params = tomllib.load(stream)
     return params
-
-
-if __name__ == "__main__":
-    import matplotlib as mpl
-    import mpl_ascii
-
-    mpl_ascii.AXES_WIDTH = 70
-    mpl_ascii.AXES_HEIGHT = 15
-
-    mpl.use("module://mpl_ascii")
-
-    filepath = "diffusion.toml"
-
-    if os.path.isfile(filepath):
-        params = load_params_from_path(filepath)
-    else:
-        params = {}
-    concentration = run_diffusion_model(**params)
-
-    np.savetxt(sys.stdout, concentration, fmt="%.6f")
diff --git a/src/diffusion_model/__main__.py b/src/diffusion_model/__main__.py
new file mode 100644
index 0000000..93fe30c
--- /dev/null
+++ b/src/diffusion_model/__main__.py
@@ -0,0 +1,23 @@
+import os
+import sys
+
+import matplotlib as mpl
+import mpl_ascii
+import numpy as np
+
+from diffusion_model import load_params_from_path, run_diffusion_model
+
+mpl_ascii.AXES_WIDTH = 70
+mpl_ascii.AXES_HEIGHT = 15
+
+mpl.use("module://mpl_ascii")
+
+filepath = "diffusion.toml"
+
+if os.path.isfile(filepath):
+    params = load_params_from_path(filepath)
+else:
+    params = {}
+concentration = run_diffusion_model(**params)
+
+np.savetxt(sys.stdout, concentration, fmt="%.6f")