aragog.output

def plot(self, num_lines: int = 11, figsize: tuple = (25, 10)) -> None:
    """Plots the solution with labelled lines according to time.

    Args:
        num_lines: Number of lines to plot. Defaults to 11.
        figsize: Size of the figure. Defaults to (25, 10).
    """
    assert self.solution is not None

    self.state.update(self.solution.y, self.solution.t)

    _, axs = plt.subplots(1, 10, sharey=True, figsize=figsize)

    # Ensure there are at least 2 lines to plot (first and last).
    num_lines = max(2, num_lines)

    # Calculate the time step based on the total number of lines.
    time_step: float = self.time_range / (num_lines - 1)

    # plot temperature
    try:
        axs[0].scatter(self.liquidus_K_staggered, self.pressure_GPa_staggered)
        axs[0].scatter(self.solidus_K_staggered, self.pressure_GPa_staggered)
    except AttributeError:
        pass

    # Plot the first line.
    def plot_times(ax, x: npt.NDArray, y: npt.NDArray) -> None:
        # If `x` is a float, create an array of the same length as `y` filled with `x`
        if np.isscalar(x):
            x = np.full((len(y), len(self.times)), x, dtype=np.float64)
        elif x.ndim == 1:
            # If `x` is 1D, reshape it or repeat it across the time dimension
            x = np.tile(x[:, np.newaxis], (1, len(self.times)))

        label_first: str = f"{self.times[0]:.2f}"
        ax.plot(x[:, 0], y, label=label_first)

        # Loop through the selected lines and plot each with a label.
        for i in range(1, num_lines - 1):
            desired_time: float = self.times[0] + i * time_step
            # Find the closest available time step.
            closest_time_index: np.intp = np.argmin(np.abs(self.times - desired_time))
            time: float = self.times[closest_time_index]
            label: str = f"{time:.2f}"  # Create a label based on the time.
            ax.plot(
                x[:, closest_time_index],
                y,
                label=label,
            )

        # Plot the last line.
        times_end: float = self.times[-1]
        label_last: str = f"{times_end:.2f}"
        ax.plot(x[:, -1], y, label=label_last)

    plot_times(axs[0], self.temperature_K_basic, self.pressure_GPa_basic)
    axs[0].set_ylabel("Pressure (GPa)")
    axs[0].set_xlabel("Temperature (K)")
    axs[0].set_title("Temperature")

    plot_times(axs[1], self.melt_fraction_staggered, self.pressure_GPa_staggered)
    axs[1].set_xlabel("Melt fraction")
    axs[1].set_title("Melt fraction")

    plot_times(axs[2], self.log10_viscosity_basic, self.pressure_GPa_basic)
    axs[2].set_xlabel("Log10(viscosity)")
    axs[2].set_title("Log10(viscosity)")

    plot_times(axs[3], self.convective_heat_flux_basic, self.pressure_GPa_basic)
    axs[3].set_xlabel("Convective heat flux")
    axs[3].set_title("Convective heat flux")

    plot_times(
        axs[4],
        self.super_adiabatic_temperature_gradient_basic,
        self.pressure_GPa_basic,
    )
    axs[4].set_xlabel("Super adiabatic temperature gradient")
    axs[4].set_title("Super adiabatic temperature gradient")

    plot_times(
        axs[5],
        self.dTdr,
        self.pressure_GPa_basic,
    )
    axs[5].set_xlabel("dTdr")
    axs[5].set_title("dTdr")

    plot_times(
        axs[6],
        self.dTdrs,
        self.pressure_GPa_basic,
    )
    axs[6].set_xlabel("dTdrs")
    axs[6].set_title("dTdrs")

    plot_times(
        axs[7],
        self.density_basic,
        self.pressure_GPa_basic,
    )
    axs[7].set_xlabel("Density")
    axs[7].set_title("Density")

    plot_times(
        axs[8],
        self.heat_capacity_basic,
        self.pressure_GPa_basic,
    )
    axs[8].set_xlabel("Heat capacity")
    axs[8].set_title("Heat capacity")

    plot_times(
        axs[9],
        self.thermal_expansivity_basic,
        self.pressure_GPa_basic,
    )
    axs[9].set_xlabel("Thermal expansivity")
    axs[9].set_title("Thermal expansivity")

    # Shrink current axis by 20% to allow space for the legend.
    # box = ax.get_position()
    # ax.set_position((box.x0, box.y0, box.width * 0.8, box.height))

    # axs[0].grid(True)

    legend = axs[2].legend()  # (loc="center left", bbox_to_anchor=(1, 0.5))
    legend.set_title("Time (yr)")
    plt.gca().invert_yaxis()

    plt.show()

def write_at_time(self, file_path: str, tidx: int = -1, compress: bool = False) -> None:
    """Write the state of the model at a particular time to a NetCDF4 file on the disk.

    Args:
        file_path: Path to the output file
        tidx: Index on the time axis at which to access the data
        compress: Whether to compress the data
    """

    logger.debug("Writing i=%d NetCDF file to %s", tidx, file_path)

    # Update the state
    assert self.solution is not None
    self.state.update(self.solution.y, self.solution.t)

    # Open the dataset
    ds: nc.Dataset = nc.Dataset(file_path, mode="w")

    # Metadata
    ds.description = "Aragog output data"
    ds.argog_version = __version__

    # Function to save scalar quantities
    def _add_scalar_variable(key: str, value: float, units: str):
        ds.createVariable(key, np.float64)
        ds[key][0] = value.item() if hasattr(value, 'item') else value
        ds[key].units = units

    # Save scalar quantities
    _add_scalar_variable("time", self.times[tidx], "yr")
    _add_scalar_variable("phi_global", self.melt_fraction_global, "")
    _add_scalar_variable("mantle_mass", self.mantle_mass, "kg")
    _add_scalar_variable("rheo_front", self.rheological_front, "")

    # Create dimensions for mesh quantities
    ds.createDimension("basic", self.shape_basic[0])
    ds.createDimension("staggered", self.shape_staggered[0])

    # Function to save mesh quantities
    def _add_mesh_variable(key: str, some_property: Any, units: str):
        if key[-2:] == "_b":
            mesh = "basic"
        elif key[-2:] == "_s":
            mesh = "staggered"
        else:
            raise KeyError(f"NetCDF variable name must end in _b or _s: {key}")
        ds.createVariable(
            key,
            np.float64,
            (mesh,),
            compression="zlib" if compress else None,
            shuffle = True,
            complevel = 4
        )
        ds[key][:] = some_property[:, tidx]
        ds[key].units = units

    # Save mesh quantities
    _add_mesh_variable("radius_b", self.radii_km_basic, "km")
    _add_mesh_variable("pres_b", self.pressure_GPa_basic, "GPa")
    _add_mesh_variable("temp_b", self.temperature_K_basic, "K")
    _add_mesh_variable("phi_b", self.melt_fraction_basic, "")
    _add_mesh_variable("Fcond_b", self.conductive_heat_flux_basic, "W m-2")
    _add_mesh_variable("Fconv_b", self.convective_heat_flux_basic, "W m-2")
    _add_mesh_variable("Fgrav_b", self.gravitational_separation_heat_flux_basic, "W m-2")
    _add_mesh_variable("Fmix_b", self.mixing_heat_flux_basic, "W m-2")
    _add_mesh_variable("Ftotal_b", self.total_heat_flux_basic, "W m-2")
    _add_mesh_variable("log10visc_b", self.log10_viscosity_basic, "Pa s")
    _add_mesh_variable("log10visc_s", self.log10_viscosity_staggered, "Pa s")
    _add_mesh_variable("density_b", self.density_basic, "kg m-3")
    _add_mesh_variable("heatcap_b", self.heat_capacity_basic, "J kg-1 K-1")
    _add_mesh_variable("mass_s", self.mass_staggered, "kg")
    _add_mesh_variable("Hradio_s", self.heating_radio, "W kg-1")
    _add_mesh_variable("Hvol_s", self.heating_dilatation, "W kg-1")
    _add_mesh_variable("Htidal_s", self.heating_tidal, "W kg-1")
    _add_mesh_variable("Htotal_s", self.heating, "W kg-1")

    # Close the dataset
    ds.close()

@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 dTdt(
    self,
    time: npt.NDArray | float,
    temperature: npt.NDArray,
) -> npt.NDArray:
    """dT/dt at the staggered nodes

    Args:
        time: Time
        temperature: Temperature at the staggered nodes

    Returns:
        dT/dt at the staggered nodes
    """
    logger.debug("temperature passed into dTdt = %s", temperature)
    # logger.debug("temperature.shape = %s", temperature.shape)
    self.state.update(temperature, time)
    heat_flux: npt.NDArray = self.state.heat_flux
    # logger.debug("heat_flux = %s", heat_flux)
    self.evaluator.boundary_conditions.apply_flux_boundary_conditions(self.state)
    # logger.debug("heat_flux = %s", heat_flux)
    # logger.debug("mesh.basic.area.shape = %s", self.data.mesh.basic.area.shape)

    energy_flux: npt.NDArray = heat_flux * self.evaluator.mesh.basic.area
    # logger.debug("energy_flux size = %s", energy_flux.shape)

    delta_energy_flux: npt.NDArray = np.diff(energy_flux, axis=0)
    # logger.debug("delta_energy_flux size = %s", delta_energy_flux.shape)
    # logger.debug("capacitance = %s", self.state.phase_staggered.capacitance.shape)
    # FIXME: Update capacitance for mixed phase (enthalpy of fusion contribution)
    capacitance: npt.NDArray = (
        self.state.capacitance_staggered() * self.evaluator.mesh.basic.volume
    )

    # Heating rate (dT/dt) from flux divergence (power per unit area)
    dTdt: npt.NDArray = -delta_energy_flux / capacitance
    logger.debug("dTdt (fluxes only) = %s", dTdt)

    # Additional heating rate (dT/dt) from internal heating (power per unit mass)
    dTdt += self.state.heating * (
        self.state.phase_staggered.density() / self.state.capacitance_staggered()
    )

    logger.debug("dTdt (with internal heating) = %s", dTdt)

    return dTdt

@classmethod
def from_file(cls, filename: str | Path, root: str | Path = Path()) -> Self:
    """Parses a configuration file

    Args:
        filename: Filename
        root: Root of the filename

    Returns:
        Parameters
    """
    configuration_file: Path = Path(root) / Path(filename)
    logger.info("Parsing configuration file = %s", configuration_file)
    parameters: Parameters = Parameters.from_file(configuration_file)

    return cls(parameters)

def initialize(self) -> None:
    """Initializes the model."""
    logger.info("Initializing %s", self.__class__.__name__)
    self.evaluator = Evaluator(self.parameters)
    self.state = State(self.parameters, self.evaluator)

def make_tsurf_event(self):
    """
    Creates a temperature event function for use with an ODE solver to monitor changes 
    in the surface temperature.The event triggers when the change exceeds the 
    threshold, allowing the solver to stop integration.

    Returns:
        The event has the attributes:
            - terminal = True: Integration stops when the event is triggered.
            - direction = -1: Only triggers when the function is decreasing through zero.
    """
    tsurf_initial = [None]

    def tsurf_event(time: float, temperature: npt.NDArray) -> float:
        """
        Event function to detect when surface temperature changes beyond a specified threshold.

        Args:
            time (float): Current time.
            temperature (np.ndarray): Current temperature profile.

        Returns:
            float: The difference between the threshold and the actual change in surface 
                temperature. When this value crosses zero from above, the event is triggered.
        """
        tsurf_current = temperature[-1] * self.parameters.scalings.temperature
        tsurf_threshold = self.parameters.solver.tsurf_poststep_change * self.parameters.scalings.temperature

        if tsurf_initial[0] is None:
            tsurf_initial[0] = tsurf_current
            return 1.0  

        delta = abs(tsurf_current - tsurf_initial[0])

        return tsurf_threshold - delta  

    tsurf_event.terminal = self.parameters.solver.event_triggering
    tsurf_event.direction = -1

    return tsurf_event

def reset(self) -> None:
    """This function initializes the model, while keeping the previous state of the
    PhaseEvaluatorCollection object. This avoids multiple loads of lookup table data
    when running Aragog multiple times.
    """
    logger.info("Resetting %s", self.__class__.__name__)
    # Update the Evaluator object except the phase properties
    self.evaluator.mesh = Mesh(self.parameters)
    self.evaluator.boundary_conditions = BoundaryConditions(self.parameters, self.evaluator.mesh)
    self.evaluator.initial_condition = InitialCondition(self.parameters, self.evaluator.mesh, self.evaluator.phases)
    # Reinstantiate the solver state
    self.state = State(self.parameters, self.evaluator)

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