aragog (package)

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()

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)