From 35fb83c8181b57f2a0213556e63ff71770631a18 Mon Sep 17 00:00:00 2001
From: Stephanie Kemna <6518317+StephanieKemna@users.noreply.github.com>
Date: Thu, 25 Jul 2024 15:31:43 +0200
Subject: [PATCH] Address some of Jorge's comments to new docstrings

---
 .../fmu_build/templates/fmu/fmu_template.cpp      |  2 +-
 .../fmu_build/templates/onnx_fmu/onnxFmu.cpp      | 15 ++++++++-------
 src/mlfmu/types/fmu_component.py                  | 11 ++++++-----
 src/mlfmu/utils/builder.py                        |  2 +-
 4 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/src/mlfmu/fmu_build/templates/fmu/fmu_template.cpp b/src/mlfmu/fmu_build/templates/fmu/fmu_template.cpp
index 94e6ffb..70eedfe 100644
--- a/src/mlfmu/fmu_build/templates/fmu/fmu_template.cpp
+++ b/src/mlfmu/fmu_build/templates/fmu/fmu_template.cpp
@@ -19,7 +19,7 @@ class
         /**
          * \brief Constructs a new {FmuName} object.
          *
-         * \param fmuResourceLocation The location of the FMU resource.
+         * \param fmuResourceLocation The location of the resources of the FMU.
          */
         {FmuName}(cppfmu::FMIString fmuResourceLocation) : OnnxFmu(fmuResourceLocation) {}
 
diff --git a/src/mlfmu/fmu_build/templates/onnx_fmu/onnxFmu.cpp b/src/mlfmu/fmu_build/templates/onnx_fmu/onnxFmu.cpp
index 77b77aa..c42794e 100644
--- a/src/mlfmu/fmu_build/templates/onnx_fmu/onnxFmu.cpp
+++ b/src/mlfmu/fmu_build/templates/onnx_fmu/onnxFmu.cpp
@@ -45,9 +45,10 @@ std::wstring OnnxFmu::formatOnnxPath(cppfmu::FMIString fmuResourceLocation)
 }
 
 /**
- * \brief Creates a session for the ONNX FMU.
+ * \brief Creates a session to the ONNX model.
  *
- * This function creates a session for the ONNX FMU using the specified ONNX model file.
+ * This function creates a session to the ONNX model, using the specified ONNX model file.
+ * This loads the weights of the model such that we can run predictions in the doStep function.
  *
  * \note The ONNX model file path must be set before calling this function.
  * \throws std::runtime_error if the session creation fails.
@@ -168,11 +169,11 @@ void OnnxFmu::GetBoolean(const cppfmu::FMIValueReference vr[], std::size_t nvr,
 }
 
 /**
- * \brief Sets the ONNX inputs for the ONNX FMU.
+ * \brief Sets the ONNX inputs for the ONNX FMU, matching FMU variables with inputs of ONNX model.
  *
- * This function sets the ONNX inputs for the ONNX FMU by iterating over the ONNX input
- * value-reference index pairs and assigning the corresponding FMU variable's real value
- * to the ONNX input.
+ * This function matches the FMU variables with the inputs to the ONNX model.
+ * It iterates over the ONNX input value-reference index pairs and assigns the corresponding FMU
+ * variable's real value to the ONNX input.
  *
  * \returns `true` if the ONNX inputs are successfully set, `false` otherwise.
  */
@@ -230,7 +231,7 @@ bool OnnxFmu::GetOnnxOutputs()
 }
 
 /**
- * \brief Initializes the ONNX states of the ONNX FMU.
+ * \brief Initializes the ONNX states of the ONNX model.
  *
  * This function initializes the ONNX states of the ONNX FMU by assigning the initial values
  * of the ONNX states from the corresponding variables in the FMU.
diff --git a/src/mlfmu/types/fmu_component.py b/src/mlfmu/types/fmu_component.py
index 3b9cc5f..87e7db7 100644
--- a/src/mlfmu/types/fmu_component.py
+++ b/src/mlfmu/types/fmu_component.py
@@ -102,9 +102,9 @@ class InternalState(BaseModelConfig):
             If this field is set, parameters for initialization will be automatically generated for these states.
         initialization_variable (Optional[str]): The name of an input or parameter in the same model interface
             that should be used to initialize this state.
-        agent_output_indexes (List[str]): Index or range of indices of agent outputs that will be stored as
+        agent_output_indexes (List[str]): Index or range of indices of agent (ONNX model) outputs that will be stored as
             internal states and will be fed as inputs in the next time step. Note: the FMU signal and the
-            agent outputs need to have the same length.
+            ONNX (agent) outputs need to have the same length.
     """
 
     name: Optional[str] = Field(
@@ -174,8 +174,8 @@ class InputVariable(Variable):
 
     Attributes
     ----------
-        agent_input_indexes (List[str]): Index or range of indices of agent inputs to which this FMU signal shall be linked to.
-            Note: The FMU signal and the agent inputs need to have the same length.
+        agent_input_indexes (List[str]): Index or range of indices of ONNX (agent) inputs to which this FMU signal shall be linked.
+            Note: The FMU signal and the ONNX (agent) inputs need to have the same length.
 
     Examples
     --------
@@ -301,7 +301,8 @@ class FmiVariable:
 
 class ModelComponent(BaseModelConfig):
     """
-    Represents a simulation model component.
+    Represents a simulation model component, used to generate the JSON schema for the model interface.
+    We define the structure of the FMU and how the inputs and outputs of the ONNX model correspond to the FMU variables.
 
     Attributes
     ----------
diff --git a/src/mlfmu/utils/builder.py b/src/mlfmu/utils/builder.py
index 955a729..bb98c69 100644
--- a/src/mlfmu/utils/builder.py
+++ b/src/mlfmu/utils/builder.py
@@ -43,7 +43,7 @@ def format_template_file(template_path: Path, save_path: Path, data: dict[str, s
 
 def create_model_description(fmu: FmiModel, src_path: Path):
     """
-    Compute XML structure for FMU and save it in a file.
+    Generate modelDescription.xml structure for FMU, and save it in a file.
 
     Args
     ----