Merge pull request #43 from crzdg/fixes-and-py311

Fixes and py311
tmcclintock · May 10, 2023 · 6e1d5ca · 6e1d5ca
2 parents f559d8f + f85e750
commit 6e1d5ca
Show file tree

Hide file tree

Showing 4 changed files with 190 additions and 165 deletions.
diff --git a/frispy/disc.py b/frispy/disc.py
@@ -1,6 +1,6 @@
 """Disc class."""
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Dict, Tuple, Type
 
 import numpy as np
@@ -63,7 +63,8 @@ class Disc:
     mass: float = 0.175  # kg
     air_density: float = 1.225  # kg / m ^ 3
     g: float = 9.81  # m / s ^ 2
-    model: Model = Model()
+    model: Model = field(default_factory=Model)
+    eom: EOM = None
     eom_class: Type = EOM
 
     def compute_trajectory(
@@ -138,6 +139,8 @@ def compute_trajectory(
             **solver_kwargs,
         )
 
+        self.eom = eom
+
         return (
             {
                 "times": result.t,

diff --git a/notebooks/Visualizing Trajectories.ipynb b/notebooks/Visualizing Trajectories.ipynb
diff --git a/pyproject.toml b/pyproject.toml
@@ -9,7 +9,6 @@ readme = "README.rst"
 python = "^3.9"
 numpy = "^1.23.5"
 scipy = "^1.9.3"
-matplotlib = "^3.6.2"
 
 [tool.poetry.group.dev.dependencies]
 isort = "^5.10.1"
@@ -22,6 +21,7 @@ sphinx-rtd-theme = "^1.1.1"
 sphinx-autoapi = "^2.0.0"
 interrogate = "^1.5.0"
 pydocstyle = "^6.1.1"
+matplotlib = "^3.6.2"
 
 [build-system]
 requires = ["poetry-core"]

diff --git a/tests/test_eom.py b/tests/test_eom.py
@@ -35,7 +35,7 @@ def test_compute_forces(self):
         for f in ["F_lift", "F_drag", "F_grav", "F_total", "Acc"]:
             assert f in result
             assert result[f].shape == (3,)
-            assert result[f].dtype == np.float
+            assert result[f].dtype == float
 
     def test_F_drag_direction(self):
         eom = EOM(**self.kwargs)
@@ -70,14 +70,14 @@ def test_compute_torques_smoke(self):
         for t in ["T_x_lab", "T_y_lab", "T_x", "T_y", "T_z", "T"]:
             assert t in result
             assert result[t].shape == (3,)
-            assert result[t].dtype == np.float
+            assert result[t].dtype == float
 
     def test_compute_derivatives_smoke(self):
         eom = EOM(**self.kwargs)
         coords = np.array([0, 0, 1, 10, 0, 0, 0, 0, 0, 0, 0, 62])
         der = eom.compute_derivatives(0, coords)
         assert der.shape == (12,)
-        assert der.dtype == np.float
+        assert der.dtype == float
 
     def test_rotation_matrix(self):
         def trig_functions(phi, theta):