Define system class

jobs/src/system.py

@@ -3,12 +3,51 @@ This module contains base class for a n-body gravitational system
 with 2 methods of simulations: direct and approximated force field.
 """
 
+from dataclasses import dataclass
 
-class gravitational_system():
-    pass
+import numpy as np
 
-    def direct_simulation():
-        pass
 
-    def force_field():
-        pass
+@dataclass
+class GravitationalSystem:
+    r0: np.ndarray[float]  # shape = (N,3)
+    v0: np.ndarray[float]  # shape = (N,3)
+    m: np.ndarray[float]  # shape = N
+    t: np.ndarray[float]  # shape = n_time_steps
+    solver: str
+
+    def __post_init__(self):
+        "Checking dimensions of inputs"
+        assert self.r0.shape == self.v0.shape
+        assert self.m.shape[0] == self.r0.shape[0]
+
+    def direct_simulation(self):
+        """
+        Using integrator to compute trajector in phase space
+        """
+        p = np.zeros((len(self.t), *self.v0.shape))
+        q = np.zeros((len(self.t), *self.r0.shape))
+        q[0] = self.r0
+        p[0] = self.m * self.v0
+
+        for i in range(1, len(self.t)):
+            # something with dt = self.t[i] - self.t[i-1]
+            # p = m*v
+            pass
+
+        return self.t, p, q
+
+    def force_field(self):
+        """
+        Using force field method to compute trajector in phase space
+        """
+        p = np.zeros((len(self.t), *self.v0.shape))
+        q = np.zeros((len(self.t), *self.r0.shape))
+        q[0] = self.r0
+        p[0] = self.m * self.v0
+        for i in range(1, len(self.t)):
+            # something with dt = self.t[i] - self.t[i-1]
+            # p = m*v
+            pass
+
+        return self.t, p, q

poetry.lock

@@ -200,6 +200,17 @@ docs = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "pygments
 testing = ["build[virtualenv]", "filelock (>=3.4.0)", "flake8-2020", "ini2toml[lite] (>=0.9)", "jaraco.develop (>=7.21)", "jaraco.envs (>=2.2)", "jaraco.path (>=3.2.0)", "pip (>=19.1)", "pytest (>=6)", "pytest-black (>=0.3.7)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=2.2)", "pytest-mypy (>=0.9.1)", "pytest-perf", "pytest-ruff", "pytest-timeout", "pytest-xdist", "tomli-w (>=1.0.0)", "virtualenv (>=13.0.0)", "wheel"]
 testing-integration = ["build[virtualenv] (>=1.0.3)", "filelock (>=3.4.0)", "jaraco.envs (>=2.2)", "jaraco.path (>=3.2.0)", "packaging (>=23.1)", "pytest", "pytest-enabler", "pytest-xdist", "tomli", "virtualenv (>=13.0.0)", "wheel"]
 
+[[package]]
+name = "toml"
+version = "0.10.2"
+description = "Python Library for Tom's Obvious, Minimal Language"
+optional = false
+python-versions = ">=2.6, !=3.0.*, !=3.1.*, !=3.2.*"
+files = [
+    {file = "toml-0.10.2-py2.py3-none-any.whl", hash = "sha256:806143ae5bfb6a3c6e736a764057db0e6a0e05e338b5630894a5f779cabb4f9b"},
+    {file = "toml-0.10.2.tar.gz", hash = "sha256:b3bda1d108d5dd99f4a20d24d9c348e91c4db7ab1b749200bded2f839ccbe68f"},
+]
+
 [[package]]
 name = "virtualenv"
 version = "20.25.0"
@@ -234,4 +245,4 @@ files = [
 [metadata]
 lock-version = "2.0"
 python-versions = "~3.11"
-content-hash = "79e1810a3524adfaaf1d9aa0a659f9dc3c9f35f46f7daa71963d06e011323054"
+content-hash = "8c0901f89e24d88032125c0e1c3297b9518cb1491f74361046bc8b771bae0fb9"

pyproject.toml

@@ -7,6 +7,7 @@ readme = "README.md"
 
 [tool.poetry.dependencies]
 python = "~3.11"
+toml = "^0.10.2"
 
 [tool.poetry.group.dev.dependencies]
 ruff = "^0.0.267"