neural_lam.custom_loggers
=========================

.. py:module:: neural_lam.custom_loggers

.. autoapi-nested-parse::

   Custom logging utilities (e.g., MLFlow wrappers) used in Neural-LAM.



Classes
-------

.. autoapisummary::

   neural_lam.custom_loggers.CustomMLFlowLogger


Module Contents
---------------

.. py:class:: CustomMLFlowLogger(experiment_name: str, tracking_uri: str, run_name: str, save_dir: str)

   Bases: :py:obj:`pytorch_lightning.loggers.MLFlowLogger`


   Custom MLFlow logger that adds the `log_image()` functionality not
   present in the default implementation from pytorch-lightning as
   of version `2.0.3` at least.

   Initialize the logger, ensure ``save_dir`` exists, and start the
   MLflow run.

   :param experiment_name: Target MLflow experiment.
   :type experiment_name: str
   :param tracking_uri: MLflow tracking server URI.
   :type tracking_uri: str
   :param run_name: Human-readable run name stored as ``mlflow.runName``.
   :type run_name: str
   :param save_dir: Directory where ``log_image`` writes temporary figure files.
                    Created eagerly with ``exist_ok=True``.
   :type save_dir: str

   .. rubric:: Notes

   Starts the MLflow run with ``log_system_metrics=True`` and also
   records ``run_id`` as an MLflow param.


   .. py:method:: log_image(key: str, images: list[matplotlib.pyplot.Figure], step: int | None = None) -> None

      Log one or more matplotlib figures as images to MLFlow.

      When ``images`` contains more than one figure, each is logged under
      a key suffixed with its position in the list (``key_0``, ``key_1``,
      ...). A single-figure list is logged under the bare key.

      :param key: Key to log the image under. If ``step`` is given, ``_{step}``
                  is appended before any per-figure index suffix.
      :type key: str
      :param images: Figures to log.
      :type images: list of matplotlib.figure.Figure
      :param step: Step to associate with the log entry. ``None`` logs without
                   a step suffix.
      :type step: int or None, optional

      :raises botocore.exceptions.NoCredentialsError: If AWS credentials for the MLflow artifact store are missing;
          re-raised for the caller to handle.



   .. py:property:: save_dir
      :type: str


      Returns the directory where the MLFlow artifacts are saved.
      Used to define the path to save output when using the logger.

      :returns: Path to the directory where the artifacts are saved.
      :rtype: str