tank.py

from abc import ABC, abstractmethod

import numpy as np

from ..mathutils.function import Function, funcify_method
from ..plots.tank_plots import _TankPlots
from ..prints.tank_prints import _TankPrints
from ..tools import tuple_handler


class Tank(ABC):
    """Abstract Tank class that defines a tank object for a rocket motor, so
    that it evaluates useful properties of the tank and its fluids, such as
    mass, volume, fluid flow rate, center of mass, etc.

    See Also
    --------
    :ref:`tanks_usage`

    Attributes
    ----------
    Tank.name : str
        Name of the tank.
    Tank.geometry : TankGeometry
        Geometry of the tank.
    Tank.flux_time : float, tuple of float, optional
        Tank flux time in seconds.
    Tank.liquid : Fluid
        Liquid inside the tank as a Fluid object.
    Tank.gas : Fluid
        Gas inside the tank as a Fluid object.
    Tank.discretize : int, optional
        Number of points to discretize fluid inputs.
    Tank.fluid_mass : Function
        Total mass of liquid and gases in kg inside the tank as a function
        of time.
    Tank.net_mass_flow_rate : Function
        Net mass flow rate of the tank in kg/s as a function of time, also
        understood as time derivative of the fluids mass.
    Tank.liquid_volume : Function
        Volume of the liquid inside the Tank in m^3 as a function of time.
    Tank.gas_volume : Function
        Volume of the gas inside the Tank in m^3 as a function of time.
    Tank.liquid_height : Function
        Height of the liquid inside the Tank in m as a function of time.
        The zero level reference is the same as set in Tank.geometry.
    Tank.gas_height : Function
        Height of the gas inside the Tank in m as a function of time.
        The zero level reference is the same as set in Tank.geometry.
    Tank.liquid_mass : Function
        Mass of the liquid inside the Tank in kg as a function of time.
    Tank.gas_mass : Function
        Mass of the gas inside the Tank in kg as a function of time.
    Tank.liquid_center_of_mass : Function
        Center of mass of the liquid inside the Tank in m as a function of
        time. The zero level reference is the same as set in Tank.geometry.
    Tank.gas_center_of_mass : Function
        Center of mass of the gas inside the Tank in m as a function of
        time. The zero level reference is the same as set in Tank.geometry.
    Tank.center_of_mass : Function
        Center of mass of liquid and gas (i.e. propellant) inside the Tank
        in m as a function of time. The zero level reference is the same as
        set in Tank.geometry.
    Tank.liquid_inertia : Function
        The inertia of the liquid inside the Tank in kg*m^2 as a function
        of time around a perpendicular axis to the Tank symmetry axis. The
        reference point is the Tank center of mass.
    Tank.gas_inertia : Function
        The inertia of the gas inside the Tank in kg*m^2 as a function of
        time around a perpendicular axis to the Tank symmetry axis. The
        reference point is the Tank center of mass.
    Tank.inertia : Function
        The inertia of the liquid and gas (i.e. propellant) inside the Tank
        in kg*m^2 as a function of time around a perpendicular axis to the
        Tank symmetry axis. The reference point is the Tank center of mass.
    """

    def __init__(self, name, geometry, flux_time, liquid, gas, discretize=100):
        """Initialize Tank class.

        Parameters
        ----------
        name : str
            Name of the tank.
        geometry : TankGeometry
            Geometry of the tank.
        flux_time : float, tuple of float, optional
            Tank flux time in seconds. It is the time range in which the tank
            flux is being analyzed. In general, during this time, the tank is
            being filled or emptied.If a float is given, the flux time is
            assumed to be between 0 and the given float, in seconds. If a tuple
            of float is given, the flux time is assumed to be between the first
            and second elements of the tuple.
        gas : Fluid
            Gas inside the tank as a Fluid object.
        liquid : Fluid
            Liquid inside the tank as a Fluid object.
        discretize : int, optional
            Number of points to discretize fluid inputs. If the input
            already has a appropriate discretization, this parameter
            must be set to None. The default is 100.
        """
        self.name = name
        self.geometry = geometry
        self.flux_time = flux_time
        self.gas = gas
        self.liquid = liquid
        self.discretize = discretize

        # Initialize plots and prints object
        self.prints = _TankPrints(self)
        self.plots = _TankPlots(self)
        return None

    @property
    def flux_time(self):
        """Returns the start and final times of the tank flux.

        Returns
        -------
        tuple
            Tuple containing start and final times of the tank flux.
        """
        return self._flux_time

    @flux_time.setter
    def flux_time(self, flux_time):
        """Sets the start and final times of the tank flux.

        Parameters
        ----------
        flux_time : tuple
            Tuple containing start and final times of the tank flux.
        """
        self._flux_time = tuple_handler(flux_time)

    @property
    @abstractmethod
    def fluid_mass(self):
        """
        Returns the total mass of liquid and gases inside the tank as a
        function of time.

        Returns
        -------
        Function
            Mass of the tank as a function of time. Units in kg.
        """
        pass

    @property
    @abstractmethod
    def net_mass_flow_rate(self):
        """
        Returns the net mass flow rate of the tank as a function of time.
        Net mass flow rate is the mass flow rate entering the tank minus the
        mass flow rate exiting the tank, including liquids and gases.

        Returns
        -------
        Function
            Net mass flow rate of the tank as a function of time.
        """
        pass

    @property
    @abstractmethod
    def fluid_volume(self):
        """
        Returns the volume total fluid volume inside the tank as a
        function of time. This volume is the sum of the liquid and gas
        volumes.

        Returns
        -------
        Function
            Volume of the fluid as a function of time.
        """
        pass

    @property
    @abstractmethod
    def liquid_volume(self):
        """
        Returns the volume of the liquid as a function of time.

        Returns
        -------
        Function
            Volume of the liquid as a function of time.
        """
        pass

    @property
    @abstractmethod
    def gas_volume(self):
        """
        Returns the volume of the gas as a function of time.

        Returns
        -------
        Function
            Volume of the gas as a function of time.
        """
        pass

    @property
    @abstractmethod
    def liquid_height(self):
        """
        Returns the liquid level as a function of time. This
        height is measured from the zero level of the tank
        geometry.

        Returns
        -------
        Function
            Height of the ullage as a function of time.
        """
        pass

    @property
    @abstractmethod
    def gas_height(self):
        """
        Returns the gas level as a function of time. This
        height is measured from the zero level of the tank
        geometry.

        Returns
        -------
        Function
            Height of the ullage as a function of time.
        """
        pass

    @property
    @abstractmethod
    def liquid_mass(self):
        """
        Returns the mass of the liquid as a function of time.

        Returns
        -------
        Function
            Mass of the liquid as a function of time.
        """
        pass

    @property
    @abstractmethod
    def gas_mass(self):
        """
        Returns the mass of the gas as a function of time.

        Returns
        -------
        Function
            Mass of the gas as a function of time.
        """
        pass

    @funcify_method("Time (s)", "Center of mass of liquid (m)")
    def liquid_center_of_mass(self):
        """
        Returns the center of mass of the liquid portion of the tank
        as a function of time. This height is measured from the zero
        level of the tank geometry.

        Returns
        -------
        Function
            Center of mass of the liquid portion of the tank as a
            function of time.
        """
        moment = self.geometry.volume_moment(
            self.geometry.bottom, self.liquid_height.max
        )
        liquid_moment = moment @ self.liquid_height
        centroid = liquid_moment / self.liquid_volume

        # Check for zero liquid volume
        bound_volume = self.liquid_volume < 1e-4 * self.geometry.total_volume
        if bound_volume.any():
            # TODO: pending Function setter impl.
            centroid.y_array[bound_volume] = self.geometry.bottom
            centroid.set_interpolation()
            centroid.set_extrapolation()

        return centroid

    @funcify_method("Time (s)", "Center of mass of gas (m)")
    def gas_center_of_mass(self):
        """
        Returns the center of mass of the gas portion of the tank
        as a function of time. This height is measured from the zero
        level of the tank geometry.

        Returns
        -------
        Function
            Center of mass of the gas portion of the tank as a
            function of time.
        """
        moment = self.geometry.volume_moment(self.geometry.bottom, self.gas_height.max)
        upper_moment = moment @ self.gas_height
        lower_moment = moment @ self.liquid_height
        centroid = (upper_moment - lower_moment) / self.gas_volume

        # Check for zero gas volume
        bound_volume = self.gas_volume < 1e-4 * self.geometry.total_volume
        if bound_volume.any():
            # TODO: pending Function setter impl.
            centroid.y_array[bound_volume] = self.liquid_height.y_array[bound_volume]
            centroid.set_interpolation()
            centroid.set_extrapolation()

        return centroid

    @funcify_method("Time (s)", "Center of mass of Fluid (m)")
    def center_of_mass(self):
        """Returns the center of mass of the tank's fluids as a function of
        time. This height is measured from the zero level of the tank
        geometry.

        Returns
        -------
        Function
            Center of mass of the tank's fluids as a function of time.
        """
        center_of_mass = (
            self.liquid_center_of_mass * self.liquid_mass
            + self.gas_center_of_mass * self.gas_mass
        ) / (self.fluid_mass)

        # Check for zero mass
        bound_mass = (
            self.fluid_mass < 0.001 * self.geometry.total_volume * self.gas.density
        )
        if bound_mass.any():
            # TODO: pending Function setter impl.
            center_of_mass.y_array[bound_mass] = self.geometry.bottom
            center_of_mass.set_interpolation()
            center_of_mass.set_extrapolation()

        return center_of_mass

    @funcify_method("Time (s)", "Inertia tensor of liquid (kg*m²)")
    def liquid_inertia(self):
        """
        Returns the inertia tensor of the liquid portion of the tank
        as a function of time. The reference point is the center of
        mass of the tank.

        Returns
        -------
        Function
            Inertia tensor of the liquid portion of the tank as a
            function of time.
        """
        Ix_volume = self.geometry.Ix_volume(
            self.geometry.bottom, self.liquid_height.max
        )
        Ix_volume = Ix_volume @ self.liquid_height

        # Steiner theorem to account for center of mass
        Ix_volume -= self.liquid_volume * self.liquid_center_of_mass**2
        Ix_volume += (
            self.liquid_volume * (self.liquid_center_of_mass - self.center_of_mass) ** 2
        )

        return self.liquid.density * Ix_volume

    @funcify_method("Time (s)", "inertia tensor of gas (kg*m^2)")
    def gas_inertia(self):
        """
        Returns the inertia tensor of the gas portion of the tank
        as a function of time. The reference point is the center of
        mass of the tank.

        Returns
        -------
        Function
            Inertia tensor of the gas portion of the tank as a
            function of time.
        """
        Ix_volume = self.geometry.Ix_volume(self.geometry.bottom, self.gas_height.max)
        lower_inertia_volume = Ix_volume @ self.liquid_height
        upper_inertia_volume = Ix_volume @ self.gas_height
        inertia_volume = upper_inertia_volume - lower_inertia_volume

        # Steiner theorem to account for center of mass
        inertia_volume -= self.gas_volume * self.gas_center_of_mass**2
        inertia_volume += (
            self.gas_volume * (self.gas_center_of_mass - self.center_of_mass) ** 2
        )

        return self.gas.density * inertia_volume

    @funcify_method("Time (s)", "inertia tensor (kg*m^2)")
    def inertia(self):
        """
        Returns the inertia tensor of the tank's fluids as a function of
        time. The reference point is the center of mass of the tank.

        Returns
        -------
        Function
            Inertia tensor of the tank's fluids as a function of time.
        """
        return self.liquid_inertia + self.gas_inertia

    def draw(self):
        """Draws the tank geometry."""
        self.plots.draw()


class MassFlowRateBasedTank(Tank):
    """Class to define a tank based on mass flow rates inputs. This class
    inherits from the Tank class. See the Tank class for more information
    on its attributes and methods.

    See Also
    --------
    :ref:`tanks_usage`
    """

    def __init__(
        self,
        name,
        geometry,
        flux_time,
        liquid,
        gas,
        initial_liquid_mass,
        initial_gas_mass,
        liquid_mass_flow_rate_in,
        gas_mass_flow_rate_in,
        liquid_mass_flow_rate_out,
        gas_mass_flow_rate_out,
        discretize=100,
    ):
        """Initializes the MassFlowRateBasedTank class.

        Parameters
        ----------
        name : str
            Name of the tank.
        geometry : TankGeometry
            Geometry of the tank.
        flux_time : float, tuple of float, optional
            Tank flux time in seconds. It is the time range in which the tank
            flux is being analyzed. In general, during this time, the tank is
            being filled or emptied.
            If a float is given, the flux time is assumed to be between 0 and
            the given float, in seconds. If a tuple of float is given, the flux
            time is assumed to be between the first and second elements of the
            tuple.
        liquid : Fluid
            Liquid inside the tank as a Fluid object.
        gas : Fluid
            Gas inside the tank as a Fluid object.
        initial_liquid_mass : float
            Initial liquid mass in the tank in kg.
        initial_gas_mass : float
            Initial gas mass in the tank in kg.
        liquid_mass_flow_rate_in : int, float, callable, string, array, Function
            Liquid mass flow rate into the tank in kg/s. Always positive.
            It must be a valid :class:`Function` source.
            If a callable is given, it must be a function of time.
            If a ``.csv`` file is given, it must have two columns, the first
            one being time in seconds and the second one being the mass flow
            rate in kg/s.
        gas_mass_flow_rate_in : int, float, callable, string, array, Function
            Gas mass flow rate into the tank in kg/s. Always positive.
            It must be a valid :class:`Function` source.
            If a callable is given, it must be a function of time.
            If a ``.csv`` file is given, it must have two columns, the first
            one being time in seconds and the second one being the mass flow
            rate in kg/s.
        liquid_mass_flow_rate_out : int, float, callable, string, array, Function
            Liquid mass flow rate out of the tank in kg/s. Always positive.
            It must be a valid :class:`Function` source.
            If a callable is given, it must be a function of time.
            If a ``.csv`` file is given, it must have two columns, the first
            one being time in seconds and the second one being the mass flow
            rate in kg/s.
        gas_mass_flow_rate_out : int, float, callable, string, array, Function
            Gas mass flow rate out of the tank in kg/s. Always positive.
            It must be a valid :class:`Function` source.
            If a callable is given, it must be a function of time.
            If a ``.csv`` file is given, it must have two columns, the first
            one being time in seconds and the second one being the mass flow
            rate in kg/s.
        discretize : int, optional
            Number of points to discretize fluid inputs. If the mass flow
            rate inputs are uniformly discretized (have the same time steps)
            this parameter may be set to None. Otherwise, an uniform
            discretization will be applied based on the discretize value.
            The default is 100.
        """
        super().__init__(name, geometry, flux_time, liquid, gas, discretize)
        self.initial_liquid_mass = initial_liquid_mass
        self.initial_gas_mass = initial_gas_mass

        # Define flow rates
        self.liquid_mass_flow_rate_in = Function(
            liquid_mass_flow_rate_in,
            inputs="Time (s)",
            outputs="Mass Flow Rate (kg/s)",
            interpolation="linear",
            extrapolation="zero",
        )
        self.gas_mass_flow_rate_in = Function(
            gas_mass_flow_rate_in,
            inputs="Time (s)",
            outputs="Mass Flow Rate (kg/s)",
            interpolation="linear",
            extrapolation="zero",
        )
        self.liquid_mass_flow_rate_out = Function(
            liquid_mass_flow_rate_out,
            inputs="Time (s)",
            outputs="Mass Flow Rate (kg/s)",
            interpolation="linear",
            extrapolation="zero",
        )
        self.gas_mass_flow_rate_out = Function(
            gas_mass_flow_rate_out,
            inputs="Time (s)",
            outputs="Mass Flow Rate (kg/s)",
            interpolation="linear",
            extrapolation="zero",
        )

        # Discretize input flow if needed
        self.discretize_flow() if discretize else None
        return None

    @funcify_method("Time (s)", "Mass (kg)")
    def fluid_mass(self):
        """
        Returns the total mass of liquid and gases inside the tank as a
        function of time.

        Returns
        -------
        Function
            Mass of the tank as a function of time. Units in kg.
        """
        return self.liquid_mass + self.gas_mass

    @funcify_method("Time (s)", "Mass (kg)")
    def liquid_mass(self):
        """
        Returns the mass of the liquid as a function of time by integrating
        the liquid mass flow rate.

        Returns
        -------
        Function
            Mass of the liquid as a function of time.
        """
        liquid_flow = self.net_liquid_flow_rate.integral_function()
        liquid_mass = self.initial_liquid_mass + liquid_flow
        if (liquid_mass < 0).any():
            raise ValueError(
                f"The tank {self.name} is underfilled. "
                + "The liquid mass is negative given the mass flow rates.\n\t\t"
                + "Try increasing the initial liquid mass, or reducing the mass"
                + "flow rates.\n\t\t"
                + f"The liquid mass is {np.min(liquid_mass.y_array):.3f} kg at "
                + f"{liquid_mass.x_array[np.argmin(liquid_mass.y_array)]} s."
            )
        return liquid_mass

    @funcify_method("Time (s)", "Mass (kg)")
    def gas_mass(self):
        """
        Returns the mass of the gas as a function of time by integrating
        the gas mass flow rate.

        Returns
        -------
        Function
            Mass of the gas as a function of time.
        """
        gas_flow = self.net_gas_flow_rate.integral_function()
        gas_mass = self.initial_gas_mass + gas_flow
        if (gas_mass < -1e-6).any():  # -1e-6 is to avoid numerical errors
            raise ValueError(
                f"The tank {self.name} is underfilled. The gas mass is negative"
                + " given the mass flow rates.\n\t\t"
                + "Try increasing the initial gas mass, or reducing the mass"
                + " flow rates.\n\t\t"
                + f"The gas mass is {np.min(gas_mass.y_array):.3f} kg at "
                + f"{gas_mass.x_array[np.argmin(gas_mass.y_array)]} s."
            )

        return gas_mass

    @funcify_method("Time (s)", "liquid mass flow rate (kg/s)", extrapolation="zero")
    def net_liquid_flow_rate(self):
        """
        Returns the net mass flow rate of liquid as a function of time.
        It is computed as the liquid mass flow rate entering the tank
        minus the liquid mass flow rate exiting the tank.

        Returns
        -------
        Function
            Net liquid mass flow rate of the tank as a function of time.
        """
        return self.liquid_mass_flow_rate_in - self.liquid_mass_flow_rate_out

    @funcify_method("Time (s)", "gas mass flow rate (kg/s)", extrapolation="zero")
    def net_gas_flow_rate(self):
        """
        Returns the net mass flow rate of gas as a function of time.
        It is computed as the gas mass flow rate entering the tank
        minus the gas mass flow rate exiting the tank.

        Returns
        -------
        Function
            Net gas mass flow rate of the tank as a function of time.
        """
        return self.gas_mass_flow_rate_in - self.gas_mass_flow_rate_out

    @funcify_method("Time (s)", "mass flow rate (kg/s)", extrapolation="zero")
    def net_mass_flow_rate(self):
        """
        Returns the net mass flow rate of the tank as a function of time.
        Net mass flow rate is the mass flow rate entering the tank minus the
        mass flow rate exiting the tank, including liquids and gases.

        Returns
        -------
        Function
            Net mass flow rate of the tank as a function of time.
        """
        return self.net_liquid_flow_rate + self.net_gas_flow_rate

    @funcify_method("Time (s)", "Volume (m³)")
    def fluid_volume(self):
        """
        Returns the volume total fluid volume inside the tank as a
        function of time. This volume is the sum of the liquid and gas
        volumes.

        Returns
        -------
        Function
            Volume of the fluid as a function of time.
        """
        return self.liquid_volume + self.gas_volume

    @funcify_method("Time (s)", "Volume (m³)")
    def liquid_volume(self):
        """
        Returns the volume of the liquid as a function of time.

        Returns
        -------
        Function
            Volume of the liquid as a function of time.
        """
        return self.liquid_mass / self.liquid.density

    @funcify_method("Time (s)", "Volume (m³)")
    def gas_volume(self):
        """
        Returns the volume of the gas as a function of time.

        Returns
        -------
        Function
            Volume of the gas as a function of time.
        """
        return self.gas_mass / self.gas.density

    @funcify_method("Time (s)", "Height (m)")
    def liquid_height(self):
        """
        Returns the liquid level as a function of time. This
        height is measured from the zero level of the tank
        geometry.

        Returns
        -------
        Function
            Height of the ullage as a function of time.
        """
        liquid_height = self.geometry.inverse_volume.compose(self.liquid_volume)
        diff_bt = liquid_height - self.geometry.bottom
        diff_up = liquid_height - self.geometry.top

        if (diff_bt < 0).any():
            raise ValueError(
                f"The tank '{self.name}' is underfilled. The liquid height is "
                + "below the tank bottom.\n\t\t"
                + "Try increasing the initial liquid mass, or reducing the mass"
                + " flow rates.\n\t\t"
                + f"The liquid height is {np.min(diff_bt.y_array):.3f} m below "
                + f"the tank bottom at {diff_bt.x_array[np.argmin(diff_bt.y_array)]:.3f} s."
            )
        if (diff_up > 0).any():
            raise ValueError(
                f"The tank '{self.name}' is overfilled. The liquid height is "
                + "above the tank top.\n\t\t"
                + "Try increasing the tank height, or reducing the initial liquid"
                + " mass, or reducing the mass flow rates.\n\t\t"
                + f"The liquid height is {np.max(diff_up.y_array):.3f} m above "
                + f"the tank top at {diff_up.x_array[np.argmax(diff_up.y_array)]:.3f} s."
            )

        return liquid_height

    @funcify_method("Time (s)", "Height (m)")
    def gas_height(self):
        """
        Returns the gas level as a function of time. This
        height is measured from the zero level of the tank
        geometry.

        Returns
        -------
        Function
            Height of the ullage as a function of time.
        """
        fluid_volume = self.gas_volume + self.liquid_volume
        gas_height = self.geometry.inverse_volume.compose(fluid_volume)
        diff = gas_height - self.geometry.top
        if (diff > 0).any():
            raise ValueError(
                f"The tank '{self.name}' is overfilled. "
                + "The gas height is above the tank top.\n\t\t"
                + "Try increasing the tank height, or reducing fluids' mass,"
                + " or double check the mass flow rates.\n\t\t"
                + f"The gas height is {np.max(diff.y_array):.3f} m above "
                + f"the tank top at {diff.x_array[np.argmax(diff.y_array)]} s."
            )
        return gas_height

    def discretize_flow(self):
        """Discretizes the mass flow rate inputs according to the flux time and
        the discretize parameter.
        """
        self.liquid_mass_flow_rate_in.set_discrete(*self.flux_time, self.discretize)
        self.gas_mass_flow_rate_in.set_discrete(*self.flux_time, self.discretize)
        self.liquid_mass_flow_rate_out.set_discrete(*self.flux_time, self.discretize)
        self.gas_mass_flow_rate_out.set_discrete(*self.flux_time, self.discretize)


class UllageBasedTank(Tank):
    """Class to define a tank whose flow is described by ullage volume, i.e.,
    the volume of the tank that is not occupied by the liquid. It assumes that
    the ullage volume is uniformly filled by the gas. This class inherits from
    the Tank class. See the Tank class for more information on its attributes
    and methods.

    See Also
    --------
    :ref:`tanks_usage`
    """

    def __init__(
        self,
        name,
        geometry,
        flux_time,
        liquid,
        gas,
        ullage,
        discretize=100,
    ):
        """
        Parameters
        ----------
        name : str
            Name of the tank.
        geometry : TankGeometry
            Geometry of the tank.
        flux_time : float, tuple of float, optional
            Tank flux time in seconds. It is the time range in which the tank
            flux is being analyzed. In general, during this time, the tank is
            being filled or emptied. If a float is given, the flux time is
            assumed to be between 0 and the given float, in seconds. If a tuple
            of float is given, the flux time is assumed to be between the first
            and second elements of the tuple.
        liquid : Fluid
            Liquid inside the tank as a Fluid object.
        gas : Fluid
            Gas inside the tank as a Fluid object.
        ullage : int, float, callable, string, array, Function
            Ullage volume as a function of time in m^3. Also understood as the
            volume of the Tank that is not occupied by liquid. Must be a valid
            :class:`Function` source.
            If a callable is given, it must be a function of time in seconds.
            If a ``.csv`` file is given, the first column must be the time in
            seconds and the second column must be the ullage volume in m^3.
        discretize : int, optional
            Number of points to discretize fluid inputs. If the ullage input is
            already discretized this parameter may be set to None. Otherwise,
            an uniform discretization will be applied based on the discretize
            value.
            The default is 100.
        """
        super().__init__(name, geometry, flux_time, liquid, gas, discretize)

        # Define ullage
        self.ullage = Function(ullage, "Time (s)", "Volume (m³)", "linear")

        # Discretize input if needed
        self.discretize_ullage() if discretize else None

        # Check if the ullage is within bounds
        if (self.ullage > self.geometry.total_volume).any():
            raise ValueError(
                "The ullage volume is out of bounds. It is greater than the "
                + "total volume of the tank."
            )
        if (self.ullage < 0).any():
            raise ValueError("The ullage volume is out of bounds. It is negative.")
        return None

    @funcify_method("Time (s)", "Mass (kg)")
    def fluid_mass(self):
        """
        Returns the total mass of liquid and gases inside the tank as a
        function of time.

        Returns
        -------
        Function
            Mass of the tank as a function of time. Units in kg.
        """
        return self.liquid_mass + self.gas_mass

    @funcify_method("Time (s)", "Mass flow rate (kg/s)")
    def net_mass_flow_rate(self):
        """
        Returns the net mass flow rate of the tank as a function of time by
        taking the derivative of the mass function.

        Returns
        -------
        Function
            Net mass flow rate of the tank as a function of time.
        """
        return self.fluid_mass.derivative_function()

    @funcify_method("Time (s)", "Volume (m³)")
    def fluid_volume(self):
        """
        Returns the volume total fluid volume inside the tank as a
        function of time. This volume is the sum of the liquid and gas
        volumes.

        Returns
        -------
        Function
            Volume of the fluid as a function of time.
        """
        return self.geometry.total_volume

    @funcify_method("Time (s)", "Volume (m³)")
    def liquid_volume(self):
        """
        Returns the volume of the liquid as a function of time. The
        volume is computed by subtracting the ullage volume from the
        total volume of the tank.

        Returns
        -------
        Function
            Volume of the liquid as a function of time.
        """
        return -(self.ullage - self.geometry.total_volume)

    @funcify_method("Time (s)", "Volume (m³)")
    def gas_volume(self):
        """
        Returns the volume of the gas as a function of time. From the
        Tank assumptions the gas volume is equal to the ullage volume.

        Returns
        -------
        Function
            Volume of the gas as a function of time.
        """
        return self.ullage

    @funcify_method("Time (s)", "Mass (kg)")
    def gas_mass(self):
        """
        Returns the mass of the gas as a function of time.

        Returns
        -------
        Function
            Mass of the gas as a function of time.
        """
        return self.gas_volume * self.gas.density

    @funcify_method("Time (s)", "Mass (kg)")
    def liquid_mass(self):
        """
        Returns the mass of the liquid as a function of time.

        Returns
        -------
        Function
            Mass of the liquid as a function of time.
        """
        return self.liquid_volume * self.liquid.density

    @funcify_method("Time (s)", "Height (m)")
    def liquid_height(self):
        """
        Returns the liquid level as a function of time. This
        height is measured from the zero level of the tank
        geometry.

        Returns
        -------
        Function
            Height of the ullage as a function of time.
        """
        return self.geometry.inverse_volume.compose(self.liquid_volume)

    @funcify_method("Time (s)", "Height (m)", "linear")
    def gas_height(self):
        """
        Returns the gas level as a function of time. This height is measured
        from the zero level of the tank geometry. Since the gas is assumed to
        be uniformly distributed in the ullage, the gas height is constant
        and equal to the top of the tank geometry.

        Returns
        -------
        Function
            Height of the ullage as a function of time.
        """
        return Function(self.geometry.top).set_discrete_based_on_model(self.gas_volume)

    def discretize_ullage(self):
        """Discretizes the ullage input according to the flux time and the
        discretize parameter."""
        self.ullage.set_discrete(*self.flux_time, self.discretize)


class LevelBasedTank(Tank):
    """Class to define a tank whose flow is described by liquid level, i.e.,
    the height of the liquid inside the tank. It assumes that the volume
    above the liquid level is uniformly occupied by gas. This class inherits
    from the Tank class. See the Tank class for more information on its
    attributes and methods.

    See Also
    --------
    :ref:`tanks_usage`
    """

    def __init__(
        self,
        name,
        geometry,
        flux_time,
        liquid,
        gas,
        liquid_height,
        discretize=100,
    ):
        """
        Parameters
        ----------
        name : str
            Name of the tank.
        geometry : TankGeometry
            Geometry of the tank.
        flux_time : float, tuple of float, optional
            Tank flux time in seconds. It is the time range in which the tank
            flux is being analyzed. In general, during this time, the tank is
            being filled or emptied. If a float is given, the flux time is
            assumed to be between 0 and the given float, in seconds. If a tuple
            of float is given, the flux time is assumed to be between the first
            and second elements of the tuple.
        liquid : Fluid
            Liquid inside the tank as a Fluid object.
        gas : Fluid
            Gas inside the tank as a Fluid object.
        liquid_height : int, float, callable, string, array, Function
            Liquid height as a function of time in m. Must be a valid
            :class:`Function` source. The liquid height zero level
            reference is assumed to be the same as the Tank geometry.
            If a callable is given, it must be a function of time in seconds
            If a ``.csv`` file is given, the first column is assumed to be the
            time and the second column the liquid height in meters.
        discretize : int, optional
            Number of points to discretize fluid inputs. If the liquid height
            input is already discretized this parameter may be set to None.
            Otherwise, an uniform discretization will be applied based on the
            discretize value.
            The default is 100.
        """
        super().__init__(name, geometry, flux_time, liquid, gas, discretize)

        # Define liquid level function
        self.liquid_level = Function(liquid_height, "Time (s)", "height (m)", "linear")

        # Discretize input if needed
        self.discretize_liquid_height() if discretize else None

        # Check if the liquid level is within bounds
        if (self.liquid_level > self.geometry.top).any():
            raise ValueError(
                "The liquid level is out of bounds. It is greater than the tank top."
            )
        if (self.liquid_level < self.geometry.bottom).any():
            raise ValueError("The liquid level is out of bounds. It is negative.")

    @funcify_method("Time (s)", "Mass (kg)")
    def fluid_mass(self):
        """
        Returns the total mass of liquid and gases inside the tank as a
        function of time.

        Returns
        -------
        Function
            Mass of the tank as a function of time. Units in kg.
        """
        # TODO: there's a bug in the net_mass_flow_rate if I don't discretize here
        sum_mass = self.liquid_mass + self.gas_mass
        sum_mass.set_discrete_based_on_model(self.liquid_level)
        return sum_mass

    @funcify_method("Time (s)", "Mass flow rate (kg/s)")
    def net_mass_flow_rate(self):
        """
        Returns the net mass flow rate of the tank as a function of time by
        taking the derivative of the mass function.

        Returns
        -------
        Function
            Net mass flow rate of the tank as a function of time.
        """
        return self.fluid_mass.derivative_function()

    @funcify_method("Time (s)", "Volume (m³)")
    def fluid_volume(self):
        """
        Returns the volume total fluid volume inside the tank as a
        function of time. This volume is the sum of the liquid and gas
        volumes.

        Returns
        -------
        Function
            Volume of the fluid as a function of time.
        """
        volume = self.gas_volume + self.liquid_volume
        diff = abs(volume - self.geometry.total_volume)
        if (diff > 1e-6).any():
            raise ValueError(
                "The `fluid_volume`, defined as the sum of `gas_volume` and "
                + "`liquid_volume`, is not equal to the total volume of the tank."
                + "\n\t\tThe difference is more than 1e-6 m^3 at "
                + f"{diff.x_array[np.argmin(diff.y_array)]} s."
            )
        return volume

    @funcify_method("Time (s)", "Volume (m³)")
    def liquid_volume(self):
        """
        Returns the volume of the liquid as a function of time.

        Returns
        -------
        Function
            Volume of the liquid as a function of time.
        """
        return self.geometry.volume.compose(self.liquid_height)

    @funcify_method("Time (s)", "Volume (m³)")
    def gas_volume(self):
        """
        Returns the volume of the gas as a function of time. The gas volume
        is assumed to uniformly occupy the volume above the liquid level.

        Returns
        -------
        Function
            Volume of the gas as a function of time.
        """
        # TODO: there's a bug on the gas_center_of_mass is I don't discretize here
        func = Function(self.geometry.total_volume)
        func -= self.liquid_volume
        func.set_discrete_based_on_model(self.liquid_volume)
        return func

    @funcify_method("Time (s)", "Height (m)")
    def liquid_height(self):
        """
        Returns the liquid level as a function of time. This height is
        measured from the zero level of the tank geometry.

        Returns
        -------
        Function
            Height of the ullage as a function of time.
        """
        return self.liquid_level

    @funcify_method("Time (s)", "Mass (kg)")
    def gas_mass(self):
        """
        Returns the mass of the gas as a function of time.

        Returns
        -------
        Function
            Mass of the gas as a function of time.
        """
        return self.gas_volume * self.gas.density

    @funcify_method("Time (s)", "Mass (kg)")
    def liquid_mass(self):
        """
        Returns the mass of the liquid as a function of time.

        Returns
        -------
        Function
            Mass of the liquid as a function of time.
        """
        return self.liquid_volume * self.liquid.density

    @funcify_method("Time (s)", "Height (m)", "linear")
    def gas_height(self):
        """
        Returns the gas level as a function of time. This
        height is measured from the zero level of the tank
        geometry. Since the gas is assumed to uniformly occupy
        the volume above the liquid level, the gas height is
        constant and equal to the top of the tank geometry.

        Returns
        -------
        Function
            Height of the ullage as a function of time.
        """
        return Function(self.geometry.top).set_discrete_based_on_model(
            self.liquid_level
        )

    def discretize_liquid_height(self):
        """Discretizes the liquid height input according to the flux time
        and the discretize parameter.
        """
        self.liquid_level.set_discrete(*self.flux_time, self.discretize)


class MassBasedTank(Tank):
    """Class to define a tank whose flow is described by liquid and gas masses.
    This class inherits from the Tank class. See the Tank class for more
    information on its attributes and methods.

    See Also
    --------
    :ref:`tanks_usage`
    """

    def __init__(
        self,
        name,
        geometry,
        flux_time,
        liquid,
        gas,
        liquid_mass,
        gas_mass,
        discretize=100,
    ):
        """
        Parameters
        ----------
        name : str
            Name of the tank.
        geometry : TankGeometry
            Geometry of the tank.
        flux_time : float, tuple of float, optional
            Tank flux time in seconds. It is the time range in which the tank
            flux is being analyzed. In general, during this time, the tank is
            being filled or emptied. If a float is given, the flux time is
            assumed to be between 0 and the given float, in seconds. If a tuple
            of float is given, the flux time is assumed to be between the first
            and second elements of the tuple.
        liquid : Fluid
            Liquid inside the tank as a Fluid object.
        gas : Fluid
            Gas inside the tank as a Fluid object.
        liquid_mass : int, float, callable, string, array, Function
            Liquid mass as a function of time in kg. Must be a valid
            :class:`Function` source.
            If a callable is given, it must be a function of time in seconds.
            If a ``.csv`` file is given, the first column must be the time in
            seconds and the second column must be the liquid mass in kg.
        gas_mass : int, float, callable, string, array, Function
            Gas mass as a function of time in kg. Must be a valid
            :class:`Function` source.
            If a callable is given, it must be a function of time in seconds.
            If a ``.csv`` file is given, the first column must be the time in
            seconds and the second column must be the gas mass in kg.
        discretize : int, optional
            Number of points to discretize fluid inputs. If the mass inputs
            are uniformly discretized (have the same time steps) this parameter
            may be set to None. Otherwise, an uniform discretization will be
            applied based on the discretize value.
            The default is 100.
        """
        super().__init__(name, geometry, flux_time, liquid, gas, discretize)

        # Define fluid masses
        self.liquid_mass = Function(liquid_mass, "Time (s)", "Mass (kg)", "linear")
        self.gas_mass = Function(gas_mass, "Time (s)", "Mass (kg)", "linear")

        # Discretize input if needed
        self.discretize_masses() if discretize else None

    @funcify_method("Time (s)", "Mass (kg)")
    def fluid_mass(self):
        """
        Returns the total mass of liquid and gases inside the tank as
        a function of time.

        Returns
        -------
        Function
            Mass of the tank as a function of time. Units in kg.
        """
        return self.liquid_mass + self.gas_mass

    @funcify_method("Time (s)", "Mass flow rate (kg/s)")
    def net_mass_flow_rate(self):
        """
        Returns the net mass flow rate of the tank as a function of time
        by taking the derivative of the mass function.

        Returns
        -------
        Function
            Net mass flow rate of the tank as a function of time.
        """
        return self.fluid_mass.derivative_function()

    @funcify_method("Time (s)", "Mass (kg)")
    def liquid_mass(self):
        """
        Returns the mass of the liquid as a function of time.

        Returns
        -------
        Function
            Mass of the liquid as a function of time.
        """
        return self.liquid_mass

    @funcify_method("Time (s)", "Mass (kg)")
    def gas_mass(self):
        """
        Returns the mass of the gas as a function of time.

        Returns
        -------
        Function
            Mass of the gas as a function of time.
        """
        return self.gas_mass

    @funcify_method("Time (s)", "Volume (m³)")
    def fluid_volume(self):
        """
        Returns the volume total fluid volume inside the tank as a
        function of time. This volume is the sum of the liquid and gas
        volumes.

        Returns
        -------
        Function
            Volume of the fluid as a function of time.
        """
        fluid_volume = self.liquid_volume + self.gas_volume

        # Check if within bounds
        diff = fluid_volume - self.geometry.total_volume

        if (diff > 1e-6).any():
            raise ValueError(
                "The `fluid_volume`, defined as the sum of `gas_volume` and "
                + "`liquid_volume`, is not equal to the total volume of the tank."
                + "\n\t\tThe difference is more than 1e-6 m^3 at "
                + f"{diff.x_array[np.argmin(diff.y_array)]} s."
            )
        
        return fluid_volume

    @funcify_method("Time (s)", "Volume (m³)")
    def gas_volume(self):
        """
        Returns the volume of the gas as a function of time.

        Returns
        -------
        Function
            Volume of the gas as a function of time.
        """
        return self.gas_mass / self.gas.density

    @funcify_method("Time (s)", "Volume (m³)")
    def liquid_volume(self):
        """
        Returns the volume of the liquid as a function of time.

        Returns
        -------
        Function
            Volume of the liquid as a function of time.
        """
        return self.liquid_mass / self.liquid.density

    @funcify_method("Time (s)", "Height (m)")
    def liquid_height(self):
        """
        Returns the liquid level as a function of time. This
        height is measured from the zero level of the tank
        geometry.

        Returns
        -------
        Function
            Height of the ullage as a function of time.
        """
        liquid_height = self.geometry.inverse_volume.compose(self.liquid_volume)
        diff_bt = liquid_height - self.geometry.bottom
        diff_up = liquid_height - self.geometry.top

        if (diff_bt < 0).any():
            raise ValueError(
                f"The tank {self.name} is underfilled. The liquid height is below "
                + "the tank bottom.\n\t\tTry increasing the initial liquid mass, "
                + "or reducing the mass flow rates.\n\t\t"
                + f"The liquid height is {np.min(diff_bt.y_array):.3f} m below "
                + f"the tank bottom at {diff_bt.x_array[np.argmin(diff_bt.y_array)]:.3f} s."
            )
        if (diff_up > 0).any():
            raise ValueError(
                f"The tank {self.name} is overfilled. The liquid height is above "
                + "the tank top.\n\t\tTry increasing the tank height, or reducing "
                + "the initial liquid mass, or reducing the mass flow rates.\n\t\t"
                + f"The liquid height is {np.max(diff_up.y_array):.3f} m above "
                + f"the tank top at {diff_up.x_array[np.argmax(diff_up.y_array)]:.3f} s."
            )

        return liquid_height

    @funcify_method("Time (s)", "Height (m)")
    def gas_height(self):
        """
        Returns the gas level as a function of time. This
        height is measured from the zero level of the tank
        geometry.

        Returns
        -------
        Function
            Height of the ullage as a function of time.
        """
        fluid_volume = self.gas_volume + self.liquid_volume
        gas_height = self.geometry.inverse_volume.compose(fluid_volume)
        diff = gas_height - self.geometry.top
        if (diff > 0).any():
            raise ValueError(
                f"The tank {self.name} is overfilled. The gas height is above "
                + "the tank top.\n\t\tTry increasing the tank height, or "
                + "reducing fluids' mass, or double check the mass flow rates."
                + f"\n\t\tThe gas height is {np.max(diff.y_array):.3f} m "
                + f"above the tank top at {diff.x_array[np.argmax(diff.y_array)]} s."
            )
        return gas_height

    def discretize_masses(self):
        """Discretizes the fluid mass inputs according to the flux time
        and the discretize parameter.
        """
        self.liquid_mass.set_discrete(*self.flux_time, self.discretize)
        self.gas_mass.set_discrete(*self.flux_time, self.discretize)