Incremental api

README.md

@@ -21,13 +21,13 @@ from essm_jax.essm import ExtendedStateSpaceModel
 tfpd = tfp.distributions
 
 
-def transition_fn(z, t, t_next):
+def transition_fn(z, t, t_next, *args):
     mean = z + jnp.sin(2 * jnp.pi * t / 10 * z)
     cov = 0.1 * jnp.eye(np.size(z))
     return tfpd.MultivariateNormalTriL(mean, jnp.linalg.cholesky(cov))
 
 
-def observation_fn(z, t):
+def observation_fn(z, t, *args):
     mean = z
     cov = t * 0.01 * jnp.eye(np.size(z))
     return tfpd.MultivariateNormalTriL(mean, jnp.linalg.cholesky(cov))

pyproject.toml

@@ -1,6 +1,32 @@
+# pyproject.toml
+
 [build-system]
-requires = [
-    "setuptools>=42",
-    "wheel"
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "essm_jax"
+version = "1.0.2"
+description = "Extended State Space Modelling in JAX"
+readme = "README.md"
+requires-python = ">3.9"
+license = { text = "Apache Software License" }
+authors = [{ name = "Joshua G. Albert", email = "[email protected]" }]
+keywords = ["kalman", "non-linear", "EKF", "modelling"]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: Apache Software License",
+    "Operating System :: OS Independent"
 ]
-build-backend = "setuptools.build_meta"
+urls = { "Homepage" = "https://github.com/joshuaalbert/essm_jax" }
+
+[project.optional-dependencies]
+# Define the extras here; they will be loaded dynamically from setup.py
+examples = []  # Placeholders; extras will load from setup.py
+
+[tool.setuptools]
+include-package-data = true
+
+
+[tool.setuptools.packages.find]
+where = ["src"]

requirements-examples.txt

@@ -1,2 +1,3 @@
 matplotlib
-jupyter
+jupyter
+jaxns

setup.py

@@ -1,36 +1,17 @@
 #!/usr/bin/env python
 
-from setuptools import find_packages
 from setuptools import setup
 
-install_requires = [
-    'jax',
-    'jaxlib',
-    'numpy<2',
-    'tensorflow_probability'
-]
 
-with open("README.md", "r") as fh:
-    long_description = fh.read()
+def load_requirements(file_name):
+    with open(file_name, "r") as file:
+        return [line.strip() for line in file if line.strip() and not line.startswith("#")]
 
-setup(name='essm_jax',
-      version='1.0.1',
-      description='Extended State Spapce Model in JAX',
-      long_description=long_description,
-      long_description_content_type="text/markdown",
-      url="https://github.com/joshuaalbert/essm_jax",
-      author='Joshua G. Albert',
-      author_email='[email protected]',
-      install_requires=install_requires,
-      tests_require=[
-          'pytest>=2.8',
-      ],
-      package_dir={'': './'},
-      packages=find_packages('./'),
-      classifiers=[
-          "Programming Language :: Python :: 3",
-          "License :: OSI Approved :: Apache Software License",
-          "Operating System :: OS Independent",
-      ],
-      python_requires='>3.9',
-      )
+
+setup(
+    install_requires=load_requirements("requirements.txt"),
+    extras_require={
+        "examples": load_requirements("requirements-examples.txt"),
+    },
+    tests_require=load_requirements("requirements-tests.txt"),
+)