TST: improve fluid tests and solve review comments.

Author: phmbressan

rocketpy/motors/fluid.py

@@ -1,13 +1,13 @@
 import numpy as np
 from scipy.constants import atm, zero_Celsius
 
-from ..mathutils.function import Function
+from ..mathutils.function import NUMERICAL_TYPES, Function
 from ..plots.fluid_plots import _FluidPlots
 from ..prints.fluid_prints import _FluidPrints
 
 
 class Fluid:
-    """Class that represents a fluid object and its simple quantities,
+    """Class that represents a fluid object and its attributes,
     such as density.
 
     Attributes
@@ -16,6 +16,7 @@ class Fluid:
         Name of the fluid.
     density : int, float, callable, string, array, Function
         Input density of the fluid in kg/m³.
+        Used internally to compute the density_function.
     density_function : Function
         Density of the fluid as a function of temperature and pressure.
     """
@@ -44,6 +45,26 @@ def __init__(self, name, density):
         self.prints = _FluidPrints(self)
         self.plots = _FluidPlots(self)
 
+    @property
+    def density(self):
+        """Density of the fluid from the class parameter input."""
+        return self._density
+
+    @density.setter
+    def density(self, value):
+        """Setter of the density class parameter."""
+        if isinstance(value, NUMERICAL_TYPES):
+            # Numerical value kept for retro-compatibility
+            self._density = value
+        else:
+            self._density = Function(
+                value,
+                interpolation="linear",
+                extrapolation="natural",
+                inputs=["Temperature (K)", "Pressure (Pa)"],
+                outputs="Density (kg/m³)",
+            )
+
     @property
     def density_function(self):
         """Density of the fluid as a function of temperature and pressure."""
@@ -60,7 +81,7 @@ def density_function(self, value):
             Density of the fluid in kg/m³ as a function of temperature
             and pressure. The value is used as a ``Function`` source.
         """
-        if isinstance(value, (int, float, np.number)):
+        if isinstance(value, NUMERICAL_TYPES):
             # pylint: disable=unused-argument
             def density_function(temperature, pressure):
                 return value
@@ -91,33 +112,35 @@ def get_time_variable_density(self, temperature, pressure):
         Function
             Density of the fluid in kg/m³ as function of time.
         """
-        if callable(temperature.source) and callable(pressure.source):
+        is_temperature_callable = callable(temperature.source)
+        is_pressure_callable = callable(pressure.source)
+        if is_temperature_callable and is_pressure_callable:
             return Function(
                 lambda time: self.density_function.get_value(
                     temperature.source(time), pressure.source(time)
                 ),
                 inputs="Time (s)",
                 outputs="Density (kg/m³)",
             )
+
+        if is_temperature_callable or is_pressure_callable:
+            time_scale = (
+                temperature.x_array if not is_temperature_callable else pressure.x_array
+            )
         else:
             time_scale = np.unique(
                 np.concatenate((temperature.x_array, pressure.x_array))
             )
-            density_time = self.density_function.get_value(
-                temperature.get_value(time_scale), pressure.get_value(time_scale)
-            )
-            return Function(
-                np.column_stack(
-                    (
-                        time_scale,
-                        density_time,
-                    )
-                ),
-                interpolation="linear",
-                extrapolation="constant",
-                inputs="Time (s)",
-                outputs="Density (kg/m³)",
-            )
+        density_time = self.density_function.get_value(
+            temperature.get_value(time_scale), pressure.get_value(time_scale)
+        )
+        return Function(
+            np.column_stack((time_scale, density_time)),
+            interpolation="linear",
+            extrapolation="constant",
+            inputs="Time (s)",
+            outputs="Density (kg/m³)",
+        )
 
     def __repr__(self):
         """Representation method.
@@ -143,15 +166,13 @@ def __str__(self):
 
     def to_dict(self, **kwargs):  # pylint: disable=unused-argument
         discretize = kwargs.get("discretize", False)
-        if not isinstance(self.density, (int, float, np.number)):
-            density = self.density_function
-        else:
-            density = self.density
 
+        density = self.density
         if discretize and isinstance(density, Function):
             density = density.set_discrete(
-                ((100, atm * 0.9), (zero_Celsius + 30, atm * 100)),
-                (25, 25),
+                lower=(100, atm * 0.9),
+                upper=(zero_Celsius + 30, atm * 50),
+                samples=(25, 25),
                 mutate_self=False,
             )
 

rocketpy/prints/fluid_prints.py

@@ -1,4 +1,4 @@
-import numpy as np
+from ..mathutils.function import NUMERICAL_TYPES
 
 
 class _FluidPrints:
@@ -36,7 +36,7 @@ def all(self):
         None
         """
         print(f"Name: {self.fluid.name}")
-        if isinstance(self.fluid.density, (int, float, np.number)):
+        if isinstance(self.fluid.density, NUMERICAL_TYPES):
             print(f"Density: {self.fluid.density:.4f} kg/m^3")
         else:
             print(f"Density: {self.fluid.density_function}")

tests/fixtures/motor/fluid_fixtures.py

@@ -4,7 +4,7 @@
 
 
 @pytest.fixture
-def nitrous_oxide_fluid():
+def nitrous_oxide_non_constant_fluid():
     """A nitrous_oxide fluid whose density varies with temperature
     and pressure.
 

tests/fixtures/motor/tanks_fixtures.py

@@ -463,14 +463,14 @@ def spherical_oxidizer_tank(oxidizer_fluid, oxidizer_pressurant):
 
 
 @pytest.fixture
-def cylindrical_variable_density_oxidizer_tank(nitrous_oxide_fluid):
+def cylindrical_variable_density_oxidizer_tank(nitrous_oxide_non_constant_fluid):
     """Fixture for creating a cylindrical variable density oxidizer
     tank. This fixture returns a `MassBasedTank` object representing
     a cylindrical oxidizer tank with variable density properties.
 
     Parameters
     ----------
-    nitrous_oxide_fluid : rocketpy.Fluid
+    nitrous_oxide_non_constant_fluid : rocketpy.Fluid
         The fluid object representing nitrous oxide.
 
     Returns
@@ -498,11 +498,11 @@ def temperature(t):
             return -40 * (t - 4) + zero_Celsius + 2
 
     return MassBasedTank(
-        name="Variable Density O2 Tank",
+        name="Variable Density N2O Tank",
         geometry=CylindricalTank(height=0.8, radius=0.06, spherical_caps=True),
         flux_time=7,
-        liquid=nitrous_oxide_fluid,
-        gas=nitrous_oxide_fluid,
+        liquid=nitrous_oxide_non_constant_fluid,
+        gas=nitrous_oxide_non_constant_fluid,
         discretize=50,
         liquid_mass=liquid_mass,
         gas_mass=0,

tests/integration/test_encoding.py

@@ -242,7 +242,7 @@ def test_rocket_encoder(rocket_name, request):
     )
 
 
-@pytest.mark.parametrize("rocket_name", ["calisto_robust"])
+@pytest.mark.parametrize("rocket_name", ["calisto_robust", "calisto_hybrid_modded"])
 def test_encoder_discretize(rocket_name, request):
     """Test encoding the total mass of ``rocketpy.Rocket`` with
     discretized encoding.

tests/unit/test_fluid.py

@@ -0,0 +1,57 @@
+import numpy as np
+from scipy.constants import atm, zero_Celsius
+
+from rocketpy.mathutils.function import Function
+from rocketpy.motors.fluid import Fluid
+
+
+def test_constant_density_fluid_properties():
+    """Test Fluid with constant density has correct properties."""
+    water = Fluid(name="Water", density=1000.0)
+    assert water.name == "Water"
+    assert water.density == 1000.0
+    assert isinstance(water.density_function, Function)
+    assert np.isclose(water.density_function.get_value(zero_Celsius, atm), 1000.0)
+
+
+def test_variable_density_fluid_properties():
+    """Test Fluid with variable density function."""
+    test_fluid = Fluid(name="TestFluid", density=lambda t, p: 44 * p / (8.314 * t))
+    assert test_fluid.name == "TestFluid"
+    assert callable(test_fluid.density)
+    assert isinstance(test_fluid.density, Function)
+    assert isinstance(test_fluid.density_function, Function)
+    expected_density = 44 * atm / (8.314 * zero_Celsius)
+    assert np.isclose(test_fluid.density_function(zero_Celsius, atm), expected_density)
+
+
+def test_get_time_variable_density_with_callable_sources():
+    """Test get_time_variable_density with callable temperature and pressure."""
+    test_fluid = Fluid("TestFluid", lambda t, p: t + p)
+    temperature = Function(lambda t: 300 + t)
+    pressure = Function(lambda t: 100_000 + 10 * t)
+    density_time_func = test_fluid.get_time_variable_density(temperature, pressure)
+    assert np.isclose(density_time_func.get_value(0), 300 + 100_000)
+    assert np.isclose(density_time_func.get_value(10), 310 + 100_100)
+
+
+def test_get_time_variable_density_with_arrays():
+    """Test get_time_variable_density with array sources."""
+    fluid = Fluid("TestFluid", lambda t, p: t + p)
+    times = np.array([0, 10])
+    temperature = Function(np.column_stack((times, [300, 310])))
+    pressure = Function(np.column_stack((times, [100_000, 100_100])))
+    density_time_func = fluid.get_time_variable_density(temperature, pressure)
+    expected = [300 + 100_000, 310 + 100_100]
+    np.testing.assert_allclose(density_time_func(times), expected)
+
+
+def test_get_time_variable_density_with_callable_and_arrays():
+    """Test get_time_variable_density with mixed sources."""
+    fluid = Fluid("TestFluid", lambda t, p: t + p)
+    times = np.array([0, 10])
+    temperature = Function(np.column_stack((times, [300, 310])))
+    pressure = Function(lambda t: 100_000 + 10 * t)
+    density_time_func = fluid.get_time_variable_density(temperature, pressure)
+    expected = [300 + 100_000, 310 + 100_100]
+    np.testing.assert_allclose(density_time_func(times), expected)