diff --git a/implicit_filter/jax_filter.py b/implicit_filter/jax_filter.py
index 2f6ed88..f9fc11a 100644
--- a/implicit_filter/jax_filter.py
+++ b/implicit_filter/jax_filter.py
@@ -7,7 +7,7 @@
 from ._jax_function import make_smooth, make_smat, make_smat_full, transform_veloctiy_to_nodes
 from ._utils import VeryStupidIdeaError, SolverNotConvergedError
 from implicit_filter.filter import Filter
-from scipy.sparse import csc_matrix, identity
+from scipy.sparse import csc_matrix, identity, spdiags
 from scipy.sparse.linalg import cg
 
 
@@ -97,6 +97,8 @@ def _compute(self, n, kl, ttu, tol=1e-6, maxiter=150000) -> np.ndarray:
         Smat1 = csc_matrix((self._ss * (1.0 / jnp.square(kl)), (self._ii, self._jj)), shape=(self._n2d, self._n2d))
         Smat = identity(self._n2d) + 0.5 * (Smat1 ** n)
 
+        # b = Smat.diagonal()
+        # pre = csc_matrix((b, (np.arange(self._n2d), np.arange(self._n2d))), shape=(self._n2d, self._n2d))
         ttw = ttu - Smat @ ttu  # Work with perturbations
 
         tts, code = cg(Smat, ttw, tol=tol, maxiter=maxiter)
diff --git a/pyproject.toml b/pyproject.toml
index b5a3c46..a36529f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -3,4 +3,42 @@ requires = [
     "setuptools>=42",
     "wheel"
 ]
-build-backend = "setuptools.build_meta"
\ No newline at end of file
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "implicit_filter"
+version = "0.1.0"
+description = "Python package implementing implicit filtering method on any type of mesh."
+readme = "readme.md"
+authors = [{ name = "Kacper Nowak", email = "kacper.nowak@awi.de" }]
+
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+]
+
+requires-python = ">=3.9,<3.11"
+
+dynamic = ["dependencies"]
+[tool.setuptools.dynamic]
+dependencies = {file = ["requirements.txt"]}
+
+[project.optional-dependencies]
+gpu = [
+    "cupy"
+]
+
+[project.urls]
+Homepage = "https://github.com/FESOM/implicit_filter"
+
+[tool.setuptools]
+platforms = ["any"]
+zip-safe  = false
+include-package-data = true
+
+[tool.setuptools.packages.find]
+where = ["implicit_filter"]
+include = ["implicit_filter*"]
+
+
diff --git a/readme.md b/readme.md
index 15dbe2e..225ecbd 100644
--- a/readme.md
+++ b/readme.md
@@ -19,6 +19,29 @@ For optimal performance usage of Nvidia GPU is highly recommended.
 
 **IO:** Xarray
 
+
+## Installation 
+Currently Python version 3.10 and 3.9 are supported. Using newer version can enabled 
+```shell
+source ./path/to/enviroment/of/your/choice
+git clone https://github.com/FESOM/implicit_filter
+
+cd implicit_filter
+# CPU only installation
+pip install -e .
+# GPU installation
+pip install -e .[gpu]
+```
+### Known issues
+Installing CuPy can cause an error, in case this happens try installing it manually:
+
+```shell
+pip install cupy
+```
+
+In case it also doesn't work, check your Nvidia driver version using `nvidia-smi` and install 
+CuPy version matching your drivers.
+
 # Tutorial
 
 Lets start with loading FESOM mesh file and data that we want to filter
@@ -45,6 +68,9 @@ from implicit_filter import CuPyFilter
 flter = CuPyFilter()
 flter.prepare_from_file(path + "fesom.mesh.diag.nc")
 ```
+JAX warning might appear about GPU not being available, but it should be ignored. 
+
+If you don't have GPU support enabled importing CuPyFilter will cause an import error.
 
 Alternatively you can set arrays by yourself, but this is shown in notebooks in examples.
 
diff --git a/requirements.txt b/requirements.txt
index d304de9..c1fedc8 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,8 +5,8 @@ contourpy==1.0.5
 cycler==0.11.0
 fastrlock==0.8
 fonttools==4.25.0
-jax==0.3.25
-jaxlib==0.3.25
+jax==0.4.25
+jaxlib==0.4.25
 kiwisolver==1.4.4
 matplotlib==3.7.1
 munkres==1.1.4
@@ -27,5 +27,4 @@ six==1.16.0
 tornado==6.3.2
 typing_extensions==4.7.1
 wheel==0.38.4
-xarray==2023.6.0
-cupy==12.1.0
\ No newline at end of file
+xarray==2023.6.0
\ No newline at end of file
diff --git a/setup.cfg b/setup.cfg
deleted file mode 100644
index b30c5fe..0000000
--- a/setup.cfg
+++ /dev/null
@@ -1,25 +0,0 @@
-[metadata]
-name = implicit_filter
-version = 0.1.0
-author = Kacper Nowak
-author_email = kacper.nowak@awi.de
-description = A small example package
-long_description = file: README.md
-long_description_content_type = text/markdown
-url = https://github.com/FESOM/implicit_filter
-project_urls =
-    Docs = https://fesom.github.io/implicit_filter/
-    Bug Tracker = https://github.com/FESOM/implicit_filter/-/issues
-    repository = https://github.com/FESOM/implicit_filter
-classifiers =
-    Programming Language :: Python :: 3
-    Operating System :: OS Independent
-
-[options]
-package_dir =
-    = implicit_filter
-packages = find:
-python_requires = >=3.9
-
-[options.packages.find]
-where = implicit_filter
\ No newline at end of file
diff --git a/setup.py b/setup.py
new file mode 100644
index 0000000..57c026b
--- /dev/null
+++ b/setup.py
@@ -0,0 +1,4 @@
+from setuptools import setup
+
+if __name__ == "__main__":
+    setup()
\ No newline at end of file