Functions for adding conditions/observables/parameter to Problem

Will simplify writing test cases and interactively assembling petab problems.

To be extended.

Related to #220.

Author: dweindl

petab/v2/problem.py

@@ -4,7 +4,9 @@
 import logging
 import os
 import tempfile
+from collections.abc import Sequence
 from math import nan
+from numbers import Number
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -724,3 +726,133 @@ def validate(
                     break
 
         return validation_results
+
+    def add_condition(self, id_: str, name: str = None, **kwargs):
+        """Add a simulation condition to the problem.
+
+        Arguments:
+            id_: The condition id
+            name: The condition name
+            kwargs: Parameter, value pairs to add to the condition table.
+        """
+        record = {CONDITION_ID: [id_], **kwargs}
+        if name is not None:
+            record[CONDITION_NAME] = name
+        tmp_df = pd.DataFrame(record).set_index([CONDITION_ID])
+        if self.condition_df is None:
+            self.condition_df = tmp_df
+        else:
+            self.condition_df = pd.concat([self.condition_df, tmp_df])
+
+    def add_observable(
+        self,
+        id_: str,
+        formula: str,
+        noise_formula: str | float | int = None,
+        noise_distribution: str = None,
+        transform: str = None,
+        name: str = None,
+        **kwargs,
+    ):
+        """Add an observable to the problem.
+
+        Arguments:
+            id_: The observable id
+            formula: The observable formula
+            noise_formula: The noise formula
+            noise_distribution: The noise distribution
+            transform: The observable transformation
+            name: The observable name
+            kwargs: additional columns/values to add to the observable table
+
+        """
+        record = {
+            OBSERVABLE_ID: [id_],
+            OBSERVABLE_FORMULA: [formula],
+        }
+        if name is not None:
+            record[OBSERVABLE_NAME] = [name]
+        if noise_formula is not None:
+            record[NOISE_FORMULA] = [noise_formula]
+        if noise_distribution is not None:
+            record[NOISE_DISTRIBUTION] = [noise_distribution]
+        if transform is not None:
+            record[OBSERVABLE_TRANSFORMATION] = [transform]
+        record.update(kwargs)
+
+        tmp_df = pd.DataFrame(record).set_index([OBSERVABLE_ID])
+        if self.observable_df is None:
+            self.observable_df = tmp_df
+        else:
+            self.observable_df = pd.concat([self.observable_df, tmp_df])
+
+    def add_parameter(
+        self,
+        id_: str,
+        estimated: bool | str | int = True,
+        nominal_value=None,
+        scale: str = None,
+        lb: Number = None,
+        ub: Number = None,
+        init_prior_type: str = None,
+        init_prior_pars: str | Sequence = None,
+        obj_prior_type: str = None,
+        obj_prior_pars: str | Sequence = None,
+        **kwargs,
+    ):
+        """Add a parameter to the problem.
+
+        Arguments:
+            id_: The parameter id
+            estimated: Whether the parameter is estimated
+            nominal_value: The nominal value of the parameter
+            scale: The parameter scale
+            lb: The lower bound of the parameter
+            ub: The upper bound of the parameter
+            init_prior_type: The type of the initialization prior distribution
+            init_prior_pars: The parameters of the initialization prior
+                distribution
+            obj_prior_type: The type of the objective prior distribution
+            obj_prior_pars: The parameters of the objective prior distribution
+            kwargs: additional columns/values to add to the parameter table
+        """
+        record = {
+            PARAMETER_ID: [id_],
+        }
+        if estimated is not None:
+            record[ESTIMATE] = [
+                int(estimated)
+                if isinstance(estimated, bool | int)
+                else estimated
+            ]
+        if nominal_value is not None:
+            record[NOMINAL_VALUE] = [nominal_value]
+        if scale is not None:
+            record[PARAMETER_SCALE] = [scale]
+        if lb is not None:
+            record[LOWER_BOUND] = [lb]
+        if ub is not None:
+            record[UPPER_BOUND] = [ub]
+        if init_prior_type is not None:
+            record[INITIALIZATION_PRIOR_TYPE] = [init_prior_type]
+        if init_prior_pars is not None:
+            if not isinstance(init_prior_pars, str):
+                init_prior_pars = PARAMETER_SEPARATOR.join(
+                    map(str, init_prior_pars)
+                )
+            record[INITIALIZATION_PRIOR_PARAMETERS] = [init_prior_pars]
+        if obj_prior_type is not None:
+            record[OBJECTIVE_PRIOR_TYPE] = [obj_prior_type]
+        if obj_prior_pars is not None:
+            if not isinstance(obj_prior_pars, str):
+                obj_prior_pars = PARAMETER_SEPARATOR.join(
+                    map(str, obj_prior_pars)
+                )
+            record[OBJECTIVE_PRIOR_PARAMETERS] = [obj_prior_pars]
+        record.update(kwargs)
+
+        tmp_df = pd.DataFrame(record).set_index([PARAMETER_ID])
+        if self.parameter_df is None:
+            self.parameter_df = tmp_df
+        else:
+            self.parameter_df = pd.concat([self.parameter_df, tmp_df])

tests/v2/test_problem.py

@@ -1,18 +1,25 @@
 import tempfile
 from pathlib import Path
 
+import numpy as np
 import pandas as pd
+from pandas.testing import assert_frame_equal
 
 import petab.v2 as petab
 from petab.v2 import Problem
 from petab.v2.C import (
     CONDITION_ID,
+    ESTIMATE,
+    LOWER_BOUND,
     MEASUREMENT,
     NOISE_FORMULA,
+    NOMINAL_VALUE,
     OBSERVABLE_FORMULA,
     OBSERVABLE_ID,
+    PARAMETER_ID,
     SIMULATION_CONDITION_ID,
     TIME,
+    UPPER_BOUND,
 )
 
 
@@ -105,3 +112,51 @@ def test_problem_from_yaml_multiple_files():
         assert petab_problem.measurement_df.shape[0] == 2
         assert petab_problem.observable_df.shape[0] == 2
         assert petab_problem.condition_df.shape[0] == 2
+
+
+def test_modify_problem():
+    """Test modifying a problem via the API."""
+    problem = Problem()
+    problem.add_condition("condition1", parameter1=1)
+    problem.add_condition("condition2", parameter2=2)
+
+    exp_condition_df = pd.DataFrame(
+        data={
+            CONDITION_ID: ["condition1", "condition2"],
+            "parameter1": [1.0, np.nan],
+            "parameter2": [np.nan, 2.0],
+        }
+    ).set_index([CONDITION_ID])
+    assert_frame_equal(
+        problem.condition_df, exp_condition_df, check_dtype=False
+    )
+
+    problem.add_observable("observable1", "1")
+    problem.add_observable("observable2", "2", noise_formula=2.2)
+
+    exp_observable_df = pd.DataFrame(
+        data={
+            OBSERVABLE_ID: ["observable1", "observable2"],
+            OBSERVABLE_FORMULA: ["1", "2"],
+            NOISE_FORMULA: [np.nan, 2.2],
+        }
+    ).set_index([OBSERVABLE_ID])
+    assert_frame_equal(
+        problem.observable_df, exp_observable_df, check_dtype=False
+    )
+
+    problem.add_parameter("parameter1", 1, 0, lb=1, ub=2)
+    problem.add_parameter("parameter2", False, 2)
+
+    exp_parameter_df = pd.DataFrame(
+        data={
+            PARAMETER_ID: ["parameter1", "parameter2"],
+            ESTIMATE: [1, 0],
+            NOMINAL_VALUE: [0.0, 2.0],
+            LOWER_BOUND: [1.0, np.nan],
+            UPPER_BOUND: [2.0, np.nan],
+        }
+    ).set_index([PARAMETER_ID])
+    assert_frame_equal(
+        problem.parameter_df, exp_parameter_df, check_dtype=False
+    )