aragog.core

def apply_flux_boundary_conditions(self, state: State) -> None:
    """Applies the boundary conditions to the state.

    Args:
        state: The state to apply the boundary conditions to
    """
    self.apply_flux_inner_boundary_condition(state)
    self.apply_flux_outer_boundary_condition(state)
    logger.debug("temperature = %s", state.temperature_basic)
    logger.debug("heat_flux = %s", state.heat_flux)

def apply_flux_inner_boundary_condition(self, state: State) -> None:
    """Applies the flux boundary condition to the state at the inner boundary.

    Args:
        state: The state to apply the boundary conditions to

    Equivalent to CORE_BC in C code.
        1: Simple core cooling
        2: Prescribed heat flux
        3: Prescribed temperature
    """
    if self._settings.inner_boundary_condition == 1:
        self.core_cooling(state)
    elif self._settings.inner_boundary_condition == 2:
        state.heat_flux[0, :] = self._settings.inner_boundary_value
    elif self._settings.inner_boundary_condition == 3:
        pass
        # raise NotImplementedError
    else:
        msg: str = (
            f"inner_boundary_condition = {self._settings.inner_boundary_condition} is unknown"
        )
        raise ValueError(msg)

def apply_flux_outer_boundary_condition(self, state: State) -> None:
    """Applies the flux boundary condition to the state at the outer boundary.

    Args:
        state: The state to apply the boundary conditions to

    Equivalent to SURFACE_BC in C code.
        1: Grey-body atmosphere
        2: Zahnle steam atmosphere
        3: Couple to atmodeller
        4: Prescribed heat flux
        5: Prescribed temperature
    """
    if self._settings.outer_boundary_condition == 1:
        self.grey_body(state)
    elif self._settings.outer_boundary_condition == 2:
        raise NotImplementedError
    elif self._settings.outer_boundary_condition == 3:
        msg: str = "Requires coupling to atmodeller"
        logger.error(msg)
        raise NotImplementedError(msg)
    elif self._settings.outer_boundary_condition == 4:
        state.heat_flux[-1, :] = self._settings.outer_boundary_value
    elif self._settings.outer_boundary_condition == 5:
        pass
    else:
        msg: str = (
            f"outer_boundary_condition = {self._settings.outer_boundary_condition} is unknown"
        )
        raise ValueError(msg)

def apply_temperature_boundary_conditions(
    self, temperature: npt.NDArray, temperature_basic: npt.NDArray, dTdr: npt.NDArray
) -> None:
    """Conforms the temperature and dTdr at the basic nodes to temperature boundary conditions.

    Args:
        temperature: Temperature at the staggered nodes
        temperature_basic: Temperature at the basic nodes
        dTdr: Temperature gradient at the basic nodes
    """
    # Core-mantle boundary
    if self._settings.inner_boundary_condition == 3:
        temperature_basic[0, :] = self._settings.inner_boundary_value
        dTdr[0, :] = (
            2 * (temperature[0, :] - temperature_basic[0, :])
            / self._mesh.basic.delta_mesh[0]
            * self._mesh.dxidr[0]
        )
    # Surface
    if self._settings.outer_boundary_condition == 5:
        temperature_basic[-1, :] = self._settings.outer_boundary_value
        dTdr[-1, :] = (
            2 * (temperature_basic[-1, :] - temperature[-1, :])
            / self._mesh.basic.delta_mesh[-1]
            * self._mesh.dxidr[-1]
        )

def apply_temperature_boundary_conditions_melt(
    self, melt_fraction: npt.NDArray, melt_fraction_basic: npt.NDArray, dphidr: npt.NDArray
) -> None:
    """Conforms the melt fraction gradient dphidr at the basic nodes
       to temperature boundary conditions.

    Args:
        melt_fraction: Melt fraction at the staggered nodes
        melt_fraction_basic: Melt fraction at the basic nodes
        dphidr: Melt fraction gradient at the basic nodes
    """
    # Core-mantle boundary
    if self._settings.inner_boundary_condition == 3:
        dphidr[0, :] = (
            2 * (melt_fraction[0, :] - melt_fraction_basic[0, :])
            / self._mesh.basic.delta_mesh[0]
            * self._mesh.dxidr[0]
        )
    # Surface
    if self._settings.outer_boundary_condition == 5:
        dphidr[-1, :] = (
            2 * (melt_fraction_basic[-1, :] - melt_fraction[-1, :])
            / self._mesh.basic.delta_mesh[-1]
            * self._mesh.dxidr[-1]
        )

def core_cooling(self, state: State) -> None:
    """Applies a core cooling heat flux according to Eq. (37) of Bower et al., 2018

    Args:
        state: The state to apply the boundary condition to
    """
    core_capacity: float = (
        4
        / 3
        * np.pi
        * np.power(self._mesh.basic.radii[0], 3)
        * self._mesh.settings.core_density
        * self._settings.core_heat_capacity
    )
    cell_capacity = self._mesh.basic.volume[0] * state.capacitance_staggered()[0, :]
    radius_ratio: float = self._mesh.basic.radii[1] / self._mesh.basic.radii[0]
    alpha = np.power(radius_ratio, 2) / ((cell_capacity / (core_capacity * 1.147)) + 1)

    state.heat_flux[0, :] = alpha * state.heat_flux[1, :]

def grey_body(self, state: State) -> None:
    """Applies a grey body flux at the surface.

    Args:
        state: The state to apply the boundary conditions to
    """
    state.heat_flux[-1, :] = (
        self._settings.emissivity
        * self._settings.scalings_.stefan_boltzmann_constant
        * (np.power(state.top_temperature, 4) - self._settings.equilibrium_temperature**4)
    )

def get_adiabat(self, pressure_basic) -> npt.NDArray:
    """Gets an adiabatic temperature profile by integrating
       the adiatiabatic temperature gradient dTdPs from the surface.
       Uses the set surface temperature.

    Args:
        Pressure field on the basic nodes

    Returns:
        Adiabatic temperature profile for the staggered nodes
    """

    def adiabat_ode(P,T):
        self._phases.active.set_pressure(P)
        self._phases.active.set_temperature(T)
        self._phases.active.update()
        return self._phases.active.dTdPs()

    # flip the pressure field top to bottom
    pressure_basic = np.flip(pressure_basic)

    sol = solve_ivp(
         adiabat_ode, (pressure_basic[0], pressure_basic[-1]),
         [self._settings.surface_temperature], t_eval=pressure_basic,
         method='RK45', rtol=1e-6, atol=1e-9)

    # flip back the temperature field from bottom to top
    temperature_basic = np.flip(sol.y[0])

    # Return temperature field at staggered nodes
    return self._mesh.quantity_at_staggered_nodes(temperature_basic)

def get_linear(self) -> npt.NDArray:
    """Gets a linear temperature profile

    Returns:
        Linear temperature profile for the staggered nodes
        Only works for uniform spatial mesh.
    """
    temperature_basic: npt.NDArray = np.linspace(
        self._settings.basal_temperature,
        self._settings.surface_temperature,
        self._mesh.basic.number_of_nodes,
    )
    return self._mesh.quantity_at_staggered_nodes(temperature_basic)

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 conductive_heat_flux(self) -> npt.NDArray:
    r"""Conductive heat flux:

    $$
    q_{cond} = -k \frac{\partial T}{\partial r}
    $$

    where $k$ is thermal conductivity, $T$ is temperature, and $r$ is radius.
    """
    conductive_heat_flux: npt.NDArray = self.phase_basic.thermal_conductivity() * -self.dTdr()

    return conductive_heat_flux

def convective_heat_flux(self) -> npt.NDArray:
    r"""Convective heat flux:

    $$
    q_{conv} = -\rho c_p \kappa_h \left( \frac{\partial T}{\partial r}
        - \left( \frac{\partial T}{\partial r} \right)_S \right)
    $$ 

    where $\rho$ is density, $c_p$ is heat capacity at constant pressure,
    $\kappa_h$ is eddy diffusivity, $T$ is temperature, $r$ is radius, and
    $S$ is entropy.
    """
    convective_heat_flux: npt.NDArray = (
        self.phase_basic.density()
        * self.phase_basic.heat_capacity()
        * self.eddy_diffusivity()
        * -self._super_adiabatic_temperature_gradient
    )

    return convective_heat_flux

def dilatation_heating(self) -> npt.NDArray:
    """Dilatation/compression heating (power per unit mass)

    Returns:
        Dilatation/compression heating (power per unit mass) at each layer of the staggered
            mesh, at a given point in time.
    """

    mass_flux_staggered: npt.NDArray = self._evaluator.mesh.quantity_at_staggered_nodes(self._mass_flux)

    dilatation_volume_source: npt.NDArray = (
        self.phase_staggered.gravitational_acceleration()
        * self.phase_staggered.delta_specific_volume()
        * mass_flux_staggered
    )

    return dilatation_volume_source

def gravitational_separation_mass_flux(self) -> npt.NDArray:
    r"""Gravitational separation mass flux:

    $$
    j_{grav} = \rho \phi (1 - \phi) v_{rel}
    $$ 

    where $\rho$ is density, $\phi$ is melt fraction, and
    $v_{rel}$ is relative velocity.
    """
    gravitational_separation_mass_flux: npt.NDArray = (
        self.phase_basic.density()
        * self.phase_basic.melt_fraction()
        * (1.0 - self.phase_basic.melt_fraction())
        * self.phase_basic.relative_velocity()
    )
    return gravitational_separation_mass_flux

def mixing_mass_flux(self) -> npt.NDArray:
    r"""Mixing mass flux:

    $$
    j_{cm} = -\rho \kappa_h \frac{\partial \phi}{\partial r}
    $$

    where $\rho$ is density, $\kappa_h$ is eddy diffusivity,
    $\phi$ is melt mass fraction, and $r$ is radius.
    """
    mixing_mass_flux: npt.NDArray = (
        self.phase_basic.density()
        * self.eddy_diffusivity()
        * -self.dphidr()
    )
    return mixing_mass_flux

def radiogenic_heating(self, time: float) -> npt.NDArray:
    """Radiogenic heating (constant with radius)

    Args:
        time: Time

    Returns:
        Radiogenic heating (power per unit mass) at each layer of the staggered
            mesh, at a given point in time.
    """

    # Total heat production at a given time (power per unit mass)
    radio_heating_float: float = 0
    for radionuclide in self._evaluator.radionuclides:
        radio_heating_float += radionuclide.get_heating(time)

    # Convert to 1D array (assuming abundances are constant)
    return radio_heating_float * np.ones_like(self.temperature_staggered)

def tidal_heating(self) -> npt.NDArray:
    """Tidal heating at each layer of the mantle.

    Args:
        time: Time

    Returns:
        Tidal heating (power per unit mass) at each layer of the staggered
            mesh, at a given point in time.
    """

    length = len(self._settings.tidal_array)

    # Must have correct shape
    if length == 1:
        # scalar
        out = np.ones_like(self.temperature_staggered) * self._settings.tidal_array[0]

    elif length == len(self.temperature_staggered):
        # vector with correct length
        out = np.array([self._settings.tidal_array]).T

    else:
        # vector with invalid length
        raise ValueError(f"Tidal heating array has invalid length {length}")

    return out

def update(self, temperature: npt.NDArray, time: FloatOrArray) -> None:
    """Updates the state.

    The evaluation order matters because we want to minimise the number of evaluations.

    Args:
        temperature: Temperature at the staggered nodes
        pressure: Pressure at the staggered nodes
        time: Time
    """
    logger.debug("Updating the state")

    logger.debug("Setting the temperature profile")
    self._temperature_staggered = temperature
    self._temperature_basic = self._evaluator.mesh.quantity_at_basic_nodes(temperature)
    logger.debug("temperature_basic = %s", self.temperature_basic)
    self._dTdr = self._evaluator.mesh.d_dr_at_basic_nodes(temperature)
    logger.debug("dTdr = %s", self.dTdr())
    self._evaluator.boundary_conditions.apply_temperature_boundary_conditions(
        temperature, self._temperature_basic, self.dTdr()
    )

    logger.debug("Setting the melt fraction profile")
    self.phase_staggered.set_temperature(temperature)
    self.phase_staggered.update()
    self.phase_basic.set_temperature(self._temperature_basic)
    self.phase_basic.update()
    phase_to_use = self._evaluator._parameters.phase_mixed.phase
    if phase_to_use == "mixed" or phase_to_use == "composite":
        self._dphidr = self._evaluator.mesh.d_dr_at_basic_nodes(
            self.phase_staggered.melt_fraction()
        )
        logger.debug("dphidr = %s", self.dphidr())
        self._evaluator.boundary_conditions.apply_temperature_boundary_conditions_melt(
            self.phase_staggered.melt_fraction(), self.phase_basic.melt_fraction(), self._dphidr
        )
    else:
        self._dphidr = np.zeros_like(self._dTdr)

    self._super_adiabatic_temperature_gradient = self.dTdr() - self.phase_basic.dTdrs()
    self._is_convective = self._super_adiabatic_temperature_gradient < 0
    velocity_prefactor: npt.NDArray = (
        self.phase_basic.gravitational_acceleration()
        * self.phase_basic.thermal_expansivity()
        * -self._super_adiabatic_temperature_gradient
    )
    # Viscous velocity
    self._viscous_velocity = (
        velocity_prefactor * self._evaluator.mesh.basic.mixing_length_cubed
    ) / (18 * self.phase_basic.kinematic_viscosity())
    self._viscous_velocity[~self.is_convective] = 0  # Must be super-adiabatic
    # Inviscid velocity
    self._inviscid_velocity = (
        velocity_prefactor * self._evaluator.mesh.basic.mixing_length_squared
    ) / 16
    self._inviscid_velocity[~self.is_convective] = 0  # Must be super-adiabatic
    self._inviscid_velocity[self._is_convective] = np.sqrt(
        self._inviscid_velocity[self._is_convective]
    )
    # Reynolds number
    self._reynolds_number = (
        self._viscous_velocity
        * self._evaluator.mesh.basic.mixing_length
        / self.phase_basic.kinematic_viscosity()
    )
    # Eddy diffusivity
    self._eddy_diffusivity = np.where(
        self.viscous_regime, self._viscous_velocity, self._inviscid_velocity
    )
    self._eddy_diffusivity *= self._evaluator.mesh.basic.mixing_length

    # Heat flux (power per unit area)
    self._heat_flux = np.zeros_like(self.temperature_basic)
    self._mass_flux = np.zeros_like(self.temperature_basic)
    if self._settings.conduction:
        self._heat_flux += self.conductive_heat_flux()
    if self._settings.convection:
        self._heat_flux += self.convective_heat_flux()
    if self._settings.gravitational_separation:
        self._mass_flux += self.gravitational_separation_mass_flux()
    if self._settings.mixing:
        self._mass_flux += self.mixing_mass_flux()
    self._heat_flux += self._mass_flux * self.phase_basic.latent_heat()

    # Heating (power per unit mass)
    self._heating = np.zeros_like(self.temperature_staggered)
    self._heating_radio = np.zeros_like(self.temperature_staggered)
    self._heating_dilatation = np.zeros_like(self.temperature_staggered)
    self._heating_tidal = np.zeros_like(self.temperature_staggered)

    if self._settings.radionuclides:
        self._heating_radio = self.radiogenic_heating(time)
        self._heating += self._heating_radio

    if self._settings.dilatation:
        self._heating_dilatation = self.dilatation_heating()
        self._heating += self._heating_dilatation

    if self._settings.tidal:
        self._heating_tidal = self.tidal_heating()
        self._heating += self._heating_tidal

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

    Args:
        scalings: scalings
    """
    self.scalings_ = scalings
    self.equilibrium_temperature /= self.scalings_.temperature
    self.core_heat_capacity /= self.scalings_.heat_capacity
    self._scale_inner_boundary_condition()
    self._scale_outer_boundary_condition()

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

    Initial condition method
        1: Linear profile
        2: User-defined temperature field (from file)
        3: Adiabatic profile

    Args:
        scalings: scalings
    """
    self.scalings_ = scalings
    self.surface_temperature /= self.scalings_.temperature
    self.basal_temperature /= self.scalings_.temperature

    if self.initial_condition == 2:
        if self.init_file == "":
            msg: str = (f"you must provide an initial temperature file")
            raise ValueError(msg)
        self.init_temperature = np.loadtxt(self.init_file)
        self.init_temperature /= self.scalings_.temperature