aragog.mesh

def get_density(self, pressure: FloatOrArray) -> npt.NDArray:
    r"""Computes density from pressure:

    $$
    \rho(P) = \rho_s \exp(P/K_S)
    $$

    where $\rho$ is density, $P$ is pressure, $\rho_s$ is surface density,
    and $K_S$ is adiabatic bulk modulus.

    Args:
        pressure: Pressure

    Returns:
        Density
    """
    density: npt.NDArray = self._surface_density * np.exp(
        pressure / self._adiabatic_bulk_modulus
    )

    return density

def get_density_from_radii(self, radii: FloatOrArray) -> FloatOrArray:
    r"""Computes density from radii:

    $$
        \rho(r) = \frac{\rho_s K_S}{K_S + \rho_s g (r-r_s)}
    $$
    where $\rho$ is density, $r$ is radius, $\rho_s$ is surface density,
    $K_S$ is adiabatic bulk modulus, and $r_s$ is surface radius.

    Args:
        radii: Radii

    Returns
        Density
    """
    density: FloatOrArray = (self._surface_density * self._adiabatic_bulk_modulus) / (
        self._adiabatic_bulk_modulus
        + self._surface_density
        * self._gravitational_acceleration
        * (radii - self._outer_boundary)
    )

    return density

def get_effective_density(self, radii) -> npt.NDArray:
    r"""
    Computes effective density using density rho(r) integration
    over a spherical shell bounded by radii

    Args:
        radii: Radii array

    Returns:
        Effective Density array
    """

    mass_shell = self.get_mass_within_shell(radii)
    volume_shell = 4 / 3 * np.pi  \
        * (np.power(radii[1:],3.0) - np.power(radii[:-1],3.0))

    return mass_shell/volume_shell

def get_mass_element(self, radii: FloatOrArray) -> npt.NDArray:
    r"""Computes the mass element:

    $$
    \frac{\delta m}{\delta r} = 4 \pi r^2 \rho
    $$

    where $\delta m$ is the mass element, $r$ is radius, and $\rho$ is
    density.

    Args:
        radii: Radii

    Returns:
        The mass element at radii
    """
    mass_element: npt.NDArray = (
        4 * np.pi * np.square(radii) * self.get_density_from_radii(radii)
    )

    return mass_element

def get_mass_within_radii(self, radii: FloatOrArray) -> npt.NDArray:
    r"""Computes mass within radii:

    $$
    m(r) = \int 4 \pi r^2 \rho dr
    $$ 
    where $m$ is mass, $r$ is radius, and $\rho$ is density.

    The integral was evaluated using WolframAlpha.

    Args:
        radii: Radii

    Returns:
        Mass within radii
    """
    a: float = self._surface_density
    b: float = self._adiabatic_bulk_modulus
    c: float = self._gravitational_acceleration
    d: float = self._outer_boundary
    beta: float = b / (a * c) - d

    def mass_integral(radii_: FloatOrArray) -> npt.NDArray:
        """Mass within radii including arbitrary constant of integration.

        Args:
            radii_: Radii

        Returns:
            Mass within radii
        """

        mass: npt.NDArray = (
            4
            * np.pi
            * b
            / c
            * (
                -1.5 * beta * beta
                - beta * radii_
                + 0.5 * radii_ * radii_
                + beta * beta * np.log(abs(beta + radii_))
            )
        )
        # + constant

        return mass

    mass: npt.NDArray = mass_integral(radii) - mass_integral(self._inner_boundary)

    return mass

def get_mass_within_shell(self, radii: npt.NDArray) -> npt.NDArray:
    """Computes the mass within spherical shells bounded by radii.

    Args:
        radii: Radii

    Returns:
        Mass within the bounded spherical shells
    """
    mass: npt.NDArray = self.get_mass_within_radii(radii[1:]) - self.get_mass_within_radii(
        radii[:-1]
    )

    return mass

def get_pressure_from_radii(self, radii: FloatOrArray) -> npt.NDArray:
    r"""Computes pressure from radii:

    $$
    P(r) = -K_S \ln \left( 1 + \frac{\rho_s g (r-r_s)}{K_S} \right)
    $$

    where $r$ is radius, $K_S$ is adiabatic bulk modulus, $P$ is pressure,
    $\rho_s$ is surface density, $g$ is gravitational acceleration, and
    $r_s$ is surface radius.

    Args:
        radii: Radii

    Returns:
        Pressure
    """
    pressure: npt.NDArray = -self._adiabatic_bulk_modulus * np.log(
        (
            self._adiabatic_bulk_modulus
            + self._surface_density
            * self._gravitational_acceleration
            * (radii - self._outer_boundary)
        )
        / self._adiabatic_bulk_modulus
    )

    return pressure

def get_pressure_gradient(self, pressure: FloatOrArray) -> npt.NDArray:
    r"""Computes the pressure gradient:

    $$
    \frac{dP}{dr} = -g \rho
    $$ 

    where $\rho$ is density, $P$ is pressure, and  $g$ is gravitational
    acceleration.

    Args:
        pressure: Pressure

    Returns:
        Pressure gradient
    """
    dPdr: npt.NDArray = -self._gravitational_acceleration * self.get_density(pressure)

    return dPdr

def get_radii_from_pressure(self, pressure: FloatOrArray) -> npt.NDArray:
    r"""Computes radii from pressure:

    $$
    P(r) = \int \frac{dP}{dr} dr = \int -g \rho_s \exp(P/K_S) dr
    $$ 

    And apply the boundary condition $P=0$ at $r=r_s$ to get:

    $$
    r(P) = \frac{K_s \left( \exp(-P/K_S)-1 \right)}{\rho_s g} + r_s
    $$

    where $r$ is radius, $K_S$ is adiabatic bulk modulus, $P$ is pressure,
    $\rho_s$ is surface density, $g$ is gravitational acceleration, and
    $r_s$ is surface radius.

    Args:
        pressure: Pressure

    Returns:
        Radii
    """
    radii: npt.NDArray = (
        self._adiabatic_bulk_modulus
        * (np.exp(-pressure / self._adiabatic_bulk_modulus) - 1)
        / (self._surface_density * self._gravitational_acceleration)
        + self._outer_boundary
    )

    return radii

def set_staggered_pressure(self, staggered_radii: npt.NDArray,) -> None:
    """Set staggered pressure based on staggered radii."""
    self._staggered_pressure = self.get_pressure_from_radii(staggered_radii)

def d_dr_at_basic_nodes(self, staggered_quantity: npt.NDArray) -> npt.NDArray:
    """Determines d/dr at the basic nodes of a quantity defined at the staggered nodes.

    Args:
        staggered_quantity: A quantity defined at the staggered nodes.

    Returns:
        d/dr at the basic nodes
    """
    d_dr_at_basic_nodes: npt.NDArray = self._d_dr_transform.dot(staggered_quantity)
    logger.debug("d_dr_at_basic_nodes = %s", d_dr_at_basic_nodes)

    return d_dr_at_basic_nodes

def get_basic_mass_coordinates_from_spatial_coordinates(self, basic_coordinates: npt.NDArray) -> npt.NDArray:
    """Computes the basic mass coordinates from basic spatial coordinates.

    Args:
        Basic spatial coordinates

    Returns:
        Basic mass coordinates
    """

    # Set the mass coordinates at the inner boundary from the core mass
    basic_mass_coordinates = np.zeros_like(basic_coordinates)
    basic_mass_coordinates[:,:] = (
        self.settings.core_density / self._planet_density
        * np.power(basic_coordinates[0,:], 3.0)
    )

    # Get mass coordinates by adding individual cell contributions to the mantle mass
    basic_volumes = (np.power(basic_coordinates[1:,:],3.0)
        - np.power(basic_coordinates[:-1,:],3.0))
    for i in range(1, self.settings.number_of_nodes):
        basic_mass_coordinates[i:,:] += (
            self.staggered_effective_density[i-1,:] * basic_volumes[i-1,:] / self._planet_density
        )

    return np.power(basic_mass_coordinates, 1.0/3.0)

def get_constant_spacing(self) -> npt.NDArray:
    """Constant radius spacing across the mantle

    Returns:
        Radii with constant spacing as a column vector
    """
    radii: npt.NDArray = np.linspace(
        self.settings.inner_radius, self.settings.outer_radius, self.settings.number_of_nodes
    )
    radii = np.atleast_2d(radii).T

    return radii

def get_dxidr_basic(self) -> npt.NDArray:
    """Computes dxidr at basic nodes."""

    dxidr = (
        self.eos.basic_density
        / self._planet_density
        * np.power(self.basic.radii,2.0)
        / np.power(self.basic.mass_radii,2.0)
    )

    return dxidr

def get_planet_density(self, basic_coordinates: npt.NDArray) -> float:
    """Computes the planet density.

    Args:
        Basic spatial coordinates

    Returns:
        Planet effective density
    """
    core_mass = self.settings.core_density *  np.power(basic_coordinates[0,0], 3.0)
    basic_volumes = (np.power(basic_coordinates[1:,0],3.0)
        - np.power(basic_coordinates[:-1,0],3.0))
    mantle_mass = np.sum(
        self.staggered_effective_density[:,0] * basic_volumes
    )
    planet_density = (
        (core_mass + mantle_mass) / (np.power(basic_coordinates[-1,0],3.0)))
    return planet_density.item()

def get_staggered_spatial_coordinates_from_mass_coordinates(self, staggered_mass_coordinates: npt.NDArray) -> npt.NDArray:
    """Computes the staggered spatial coordinates from staggered mass coordinates.

    Args:
        Staggered mass coordinates

    Returns:
        Staggered spatial coordinates
    """

    # Initialise the staggered spatial coordinate to the inner boundary
    staggered_coordinates = (np.ones_like(staggered_mass_coordinates)
        * np.power(self.settings.inner_radius,3.0))

    # Add first half cell contribution
    staggered_coordinates += (
        self._planet_density
         * (np.power(staggered_mass_coordinates[0,:],3.0)
        - np.power(self.basic.mass_radii[0,:], 3.0))
        / self.staggered_effective_density[0,:]
    )

    # Get spatial coordinates by adding individual cell contributions to the mantle mass
    shell_effective_density = 0.5*(
        self.staggered_effective_density[1:,:] + self.staggered_effective_density[:-1,:])
    shell_mass_volumes = (np.power(staggered_mass_coordinates[1:,:],3.0)
        - np.power(staggered_mass_coordinates[:-1,:],3.0))
    for i in range(1,self.settings.number_of_nodes-1):
        staggered_coordinates[i:,:] += (
            self._planet_density
            * shell_mass_volumes[i-1,:]
            / shell_effective_density[i-1,:]
        )

    return np.power(staggered_coordinates, 1.0/3.0)

def quantity_at_basic_nodes(self, staggered_quantity: npt.NDArray) -> npt.NDArray:
    """Determines a quantity at the basic nodes that is defined at the staggered nodes.

    Uses backward and forward differences at the inner and outer radius, respectively, to
    obtain the quantity values of the basic nodes at the innermost and outermost nodes.
    When using temperature boundary conditions, values at outer boundaries will be overwritten.
    When using flux boundary conditions, values at outer boundaries will be used to provide
    estimate of individual components of heat fluxes though the total heat flux is imposed.

    Args:
        staggered_quantity: A quantity defined at the staggered nodes

    Returns:
        The quantity at the basic nodes
    """
    quantity_at_basic_nodes: npt.NDArray = self._quantity_transform.dot(staggered_quantity)
    logger.debug("quantity_at_basic_nodes = %s", quantity_at_basic_nodes)

    return quantity_at_basic_nodes

def quantity_at_staggered_nodes(self, basic_quantity: npt.NDArray) -> npt.NDArray:
    """Determines a quantity at the staggered nodes that is defined at the basic nodes.

    Staggered nodes are always located at cell centers, whatever the mesh.

    Args:
        basic_quantity: A quantity defined at the basic nodes

    Returns:
        The quantity at the staggered nodes
    """
    quantity_at_staggered_nodes: npt.NDArray = 0.5 * (
        basic_quantity[:-1, ...] + basic_quantity[1:, ...])
    logger.debug("quantity_at_staggered_nodes = %s", quantity_at_staggered_nodes)

    return quantity_at_staggered_nodes

@classmethod
def from_file(cls, *filenames) -> Self:
    """Parses the parameters in a configuration file(s)

    Args:
        *filenames: Filenames of the configuration data
    """
    parser: ConfigParser = ConfigParser()
    parser.read(*filenames)

    init_dict: dict[str, Any] = {}
    for section_name, dataclass_ in _get_dataclass_from_section_name().items():
        init_dict[section_name] = parser.parse_section(
            using_dataclass=dataclass_, section_name=section_name
        )
    radionuclides: list[_Radionuclide] = []
    for radionuclide_section in cls.radionuclide_sections(parser):
        radionuclide = parser.parse_section(
            using_dataclass=_Radionuclide, section_name=radionuclide_section
        )
        radionuclides.append(radionuclide)

    init_dict["radionuclides"] = radionuclides

    return cls(**init_dict)  # Unpacking gives required arguments so pylint: disable=E1125

@staticmethod
def radionuclide_sections(parser: ConfigParser) -> list[str]:
    """Section names relating to radionuclides

    Sections relating to radionuclides must have the prefix radionuclide_
    """
    return [
        parser[section].name
        for section in parser.sections()
        if section.startswith("radionuclide_")
    ]

def set_staggered_pressure(self, staggered_radii: npt.NDArray,) -> None:
    """Set staggered pressure based on staggered radii."""
    self._staggered_pressure = self._interp_pressure(staggered_radii).reshape(-1,1)

def scale_attributes(self, scalings: _ScalingsParameters) -> None:
    """Scales the attributes

    Args:
        scalings: scalings
    """
    self.scalings_ = scalings
    self.outer_radius /= self.scalings_.radius
    self.inner_radius /= self.scalings_.radius
    self.core_density /= self.scalings_.density
    self.surface_density /= self.scalings_.density
    self.gravitational_acceleration /= self.scalings_.gravitational_acceleration
    self.adiabatic_bulk_modulus /= self.scalings_.pressure

    if self.eos_method == 2:
        if self.eos_file == "":
            msg: str = (f"you must provide a file for setting up equation of state")
            raise ValueError(msg)
        arr = np.loadtxt(self.eos_file)
        self.eos_radius = arr[:,0] / self.scalings_.radius
        self.eos_pressure = arr[:,1] / self.scalings_.pressure
        self.eos_density = arr[:,2] / self.scalings_.density
        self.eos_gravity = arr[:,3] / self.scalings_.gravitational_acceleration
        # Check that provided eos radius roughly match with Aragog mesh
        if ((self.eos_radius[0] < self.inner_radius) or
            (self.eos_radius[-1] > self.outer_radius) or
            (self.eos_radius[-1]-self.eos_radius[0]) < 0.75*(self.outer_radius-self.inner_radius)):
            msg: str = (f"Radius array in EOS file: Values out of range.")
            raise ValueError(msg)