ml_coupler.py

Go to the documentation of this file.

#!/usr/bin/env python3
 
# Copyright (c) 2026 The YAC Authors
#
# SPDX-License-Identifier: BSD-3-Clause
 
import numpy as np
from yac import Action, ExchangeType, Field, Location, Reg2dGrid, TimeUnit, YAC
 
 
class MLCoupler:
    """
    Encapsulates YAC setup and exchange for the ml_predictor component.
    """
 
    def __init__(self, component_name: str = "ml_predictor"):
        """
        Initialize the YAC coupler for the ml_predictor component.
 
        @param[in] component_name Name of the component for registration in YAC
        """
        self.component_name = component_name
        self.yac = YAC()
        self.comp = self.yac.def_comp(self.component_name)
        self.temp_field = None
        self.heat_flux_field = None
        if self.comp.comp_comm.size > 1:
            raise RuntimeError("ml_predictor must run in a single process")
 
    def setup(self, grid_name: str, x_vertices: np.ndarray, y_vertices: np.ndarray):
        """
        Define the 2D regular grid and fields for the ml_predictor component.
 
        @param[in] grid_name Name of the grid for registration in YAC
        @param[in] x_vertices X coordinates of the grid vertices
        @param[in] y_vertices Y coordinates of the grid vertices
        """
        grid = Reg2dGrid(grid_name, x_vertices, y_vertices,
                         cyclic=[False, False])
 
        # Compute the cell centers from the vertex coordinates
        x_centers = 0.5 * (x_vertices[:-1] + x_vertices[1:])
        y_centers = 0.5 * (y_vertices[:-1] + y_vertices[1:])
 
        # Define the points for the grid cells where the fields will be defined
        points = grid.def_points(Location.CELL, x_centers, y_centers)
 
        # Define the fields for temperature input and heat flux prediction
        self.temp_field = Field.create(
            "temperature", self.comp, points,
            collection_size=1, timestep="1", timeunit=TimeUnit.SECOND)
 
        self.heat_flux_field = Field.create(
            "heat_flux", self.comp, points,
            collection_size=1, timestep="1", timeunit=TimeUnit.SECOND)
 
        # Finalize the component definition
        # This is a collective operation and must be called by all processes.
        # It computes the interpolation weights and sets up communication
        # patterns for the exchange of data between the coupled components.
        self.yac.enddef()
 
    def has_data(self) -> bool:
        return self.temp_field.action != Action.OUT_OF_BOUND
 
    @property
    def heat_flux_is_coupled(self) -> bool:
        """
        Whether heat_flux is actually coupled back to the simulation
        component. A one-way (training) configuration only couples
        temperature, in which case heat_flux's role is
        ExchangeType.NONE instead of ExchangeType.SOURCE.
 
        @return True if heat_flux should be sent via `send_heat_flux`
        """
        return self.heat_flux_field.role == ExchangeType.SOURCE
 
    @property
    def datetime(self) -> str:
        """
        Return the current datetime of the next exchange
           (temperature is used as a reference).
 
        @return The current datetime of the next exchange as an ISO 8601
                string (e.g., "2024-01-01T00:00:00Z")
        """
        return self.temp_field.datetime
 
    def recv_temperature(self) -> np.ndarray | None:
        """
        Receive the temperature field from the Fortran component.
 
        @return The temperature values, or None if no data was received.
        """
        temperature, info = self.temp_field.get()
        assert info in (Action.COUPLING, Action.GET_FOR_RESTART, Action.NONE), (
            f"Unexpected action: {info}"
        )
        if temperature is None:
            return None
        return temperature.ravel()
 
    def send_heat_flux(self, heat_flux: np.ndarray):
        """
        Send the predicted heat flux back to the Fortran component.
 
        Only call this if `heat_flux_is_coupled` is True; otherwise there is
        no coupling configured for heat_flux to send data on (e.g. during a
        one-way training run).
 
        @param[in] heat_flux The predicted heat flux values to send back.
        """
        assert self.heat_flux_is_coupledheat_flux_is_coupled, (
            "heat_flux is not coupled in the current configuration; "
            "check heat_flux_is_coupled before calling send_heat_flux"
        )
        # Send the predicted heat flux back to the Fortran component
        self.heat_flux_field.put(heat_flux.reshape(1, -1))
 
    def finalize(self):
        """Finalize the YAC coupler and clean up resources."""
        if self.yac is not None:
            self.yac.cleanup()
            self.yac = None