Repository: neurodiffgym/neurodiffeq
Issue: BundleSolver setup too restrictive

The following code block assumes each bundle parameter must fall into one of the following categories.
1. a parameter that only appears in the boundary/initial condition **(NOT in the equation)**; or
2. a parameter that only appears in the equation **(NOT in the boundary/initial condition)**
https://github.com/NeuroDiffGym/neurodiffeq/blob/a3df9fd221223a85ebc973002b8e3d2b280cbc19/neurodiffeq/solvers.py#L1318-L1332

However, in theory, there should be a case where the parameter appears both in the equation and in the initial/boundary condition. The `var` here probably implies a parameter that appears only in the equation. A fix is needed.

Changed files:
neurodiffeq/generators.py, tests/test_conditions.py, neurodiffeq/conditions.py, neurodiffeq/solvers.py

diff --git a/neurodiffeq/conditions.py b/neurodiffeq/conditions.py
index 3fafef5..397803a 100644
--- a/neurodiffeq/conditions.py
+++ b/neurodiffeq/conditions.py
@@ -75,6 +75,66 @@ class BaseCondition:
         self.ith_unit = ith_unit
 
 
+class _BundleConditionMixin:
+    def __init__(self, bundle_param_lookup=None, allowed_params=None):
+        """Mixin class for bundle conditions. Unlike other conditions, parameters of a bundle conditions are not fixed.
+        For example, a regular `neurodiffeq.conditions.IVP` instance only allows fixed t_0 and u_0. But a bundle
+        condition allows for t_0 and/or u_0 to be dynamically sampled by generators during training. One can use this to
+        learn a bundle of solutions (as a function of t_0 and/or u_0).
+
+        :param bundle_param_lookup:
+            A dictionary that maps bundle parameter name to its corresponding index.
+            For example, if bundle_param_lookup == {'t_0': 0, 'u_0': 1}, then the `thetas` passed to self.enforce()`
+            must contain two entries, with thetas[0] being a tensor of `t_0` and thetas[1] being a tensor of `u_0`.
+            Similarly, if bundle_param_lookup == {'t_0': 1, 'u_0': 0}, then the two entries in thetas should be
+            swapped. Defaults to empty dictionary.
+        :type bundle_param_lookup: Dict[str, int]
+        :param allowed_params:
+            A collection of parameter names allowed in ``bundle_param_lookup``. If specified, an error will be raised
+            if ``bundle_param_lookup`` contains names that do not appear in ``allowed_params``.
+        :type allowed_params: Set[str] or List[str] or Tuple[str]
+        """
+        self.bundle_param_lookup = bundle_param_lookup or {}
+
+        if isinstance(allowed_params, str):
+            allowed_params = set(allowed_params)
+
+        if allowed_params:
+            illegal_params = set(self.bundle_param_lookup) - set(allowed_params)
+            if illegal_params:
+                raise ValueError(
+                    f"The following parameter(s) are not allowed in `bundle_parameters_lookup`: {illegal_params}.\\n"
+                    f"Supported parameter name(s) are: {allowed_params}.")
+
+    def _get_parameter(self, param_name, thetas, override_name=None):
+        """Private method to get a parameter by name (e.g. `t_0`).
+        If the parameter is present in `self.bundle conditions`, return the corresponding value from `thetas`.
+        Otherwise, return the default value (evaluated during instantiation).
+
+        Example:
+        - If param_name == 't_0', and self.bundle_param_lookup['t_0'] == 0, the method will return `thetas[0]`.
+        - If `t_0` is not present in self.bundle_param_lookup, it will return `self.t_0`, or, self.some_other_attribute
+          if override_name is set to 'some_other_attribute'.
+
+        :param param_name: Name of parameter to be used. E.g. `t_0`, `u_0` in the case of an initial value problem.
+        :type param_name: str
+        :param thetas:
+            Sequence of Bundle parameters. The order must follow what is specified in `self.bundle_param_lookup'.
+            For example, if self.bundle_param_lookup == {'t_0': 0, 'u_0': 1}, then thetas[0] must be `t_0`
+            and thetas[1] must be `u_0`.
+        :type thetas: Tuple[torch.Tensor]
+        :param override_name:
+            Name to override `param_name` when `param_name` is not contained in `self.bundle_param_lookup`.
+            See example above to see how it is used.
+        :type override_name: str
+        :return: The parameter value specified by the bundle or its default value (evaluated during instantiation)
+        :rtype: int or float or torch.Tensor
+        """
+        if param_name in self.bundle_param_lookup:
+            return thetas[self.bundle_param_lookup[param_name]]
+        return getattr(self, override_name or param_name)
+
+
 class IrregularBoundaryCondition(BaseCondition):
     # Is there a more elegant solution?
     def in_domain(self, *coordinates):
@@ -207,7 +267,7 @@ class IVP(BaseCondition):
             return self.u_0 + (t - self.t_0) * self.u_0_prime + ((1 - torch.exp(-t + self.t_0)) ** 2) * output_tensor
 
 
-class BundleIVP(BaseCondition):
+class BundleIVP(BaseCondition, _BundleConditionMixin):
     r"""An initial value problem of one of the following forms:
 
     - Dirichlet condition: :math:`u(t_0,\\boldsymbol{\\theta})=u_0`.
@@ -216,31 +276,31 @@ class BundleIVP(BaseCondition):
     Here :math:`\\boldsymbol{\\theta}=(\\theta_{1},\\theta_{2},...,\\theta_{n})\\in\\mathbb{R}^n`,
     where each :math:`\\theta_i` represents a parameter, or a condition, of the ODE system that we want to solve.
 
-    :param t_0: The initial time.
+    :param t_0: The initial time. Ignored if present in `bundle_param_lookup`.
     :type t_0: float
-    :param u_0: The initial value of :math:`u`. :math:`u(t_0,\\boldsymbol{\\theta})=u_0`.
+    :param u_0:
+        The initial value of :math:`u`. :math:`u(t_0,\\boldsymbol{\\theta})=u_0`.
+        Ignored if present in `bundle_param_lookup`.
     :type u_0: float
     :param u_0_prime:
         The initial derivative of :math:`u` w.r.t. :math:`t`.
         :math:`\\displaystyle\\frac{\\partial u}{\\partial t}\\bigg|_{t = t_0}(\\boldsymbol{\\theta}) = u_0'`.
         Defaults to None.
+        Ignored if present in `bundle_param_lookup`.
     :type u_0_prime: float, optional
-    :param bundle_conditions:
-        The initial conditions that will be included in the total bundle,
-        in addition to the parameters of the ODE system.
-        The values asociated with their respective keys used in bundle_conditions
-        (e.g bundle_conditions={'t_0': 0, 'u_0': 1, 'u_0_prime': 2}),
-        must reflect the index of the tuple used in theta_min and theta_max in ``neurodiffeq.solvers.BundleSolver1D``,
-        (e.g theta_min=(t_0_min, u_0_min, u_0_prime_min)).
-        Defaults to {}
-    :type bundle_conditions: dict{str: int, ..., str: int}
+    :param bundle_param_lookup: See _BundleConditionMixin for details. Allowed keys are 't_0', 'u_0', and 'u_0_prime'.
+    :type bundle_param_lookup: Dict[str, int]
     """
 
-    @deprecated_alias(x_0='u_0', x_0_prime='u_0_prime')
-    def __init__(self, t_0=None, u_0=None, u_0_prime=None, bundle_conditions={}):
-        super().__init__()
+    @deprecated_alias(x_0='u_0', x_0_prime='u_0_prime', bundle_conditions='bundle_param_lookup')
+    def __init__(self, t_0=None, u_0=None, u_0_prime=None, bundle_param_lookup=None):
+        BaseCondition.__init__(self)
+        _BundleConditionMixin.__init__(
+            self,
+            bundle_param_lookup=bundle_param_lookup,
+            allowed_params=['t_0', 'u_0', 'u_0_prime'],
+        )
         self.t_0, self.u_0, self.u_0_prime = t_0, u_0, u_0_prime
-        self.bundle_conditions = bundle_conditions
 
     def parameterize(self, output_tensor, t, *theta):
         r"""Re-parameterizes outputs such that the Dirichlet/Neumann condition is satisfied.
@@ -275,25 +335,64 @@ class BundleIVP(BaseCondition):
         :rtype: `torch.Tensor`
         """
 
-        t_0, u_0, u_0_prime = self.t_0, self.u_0, self.u_0_prime
-        if 'u_0' in self.bundle_conditions:
-            u_0 = theta[self.bundle_conditions['u_0']]
-        if 'u_0_prime' in self.bundle_conditions:
-            u_0_prime = theta[self.bundle_conditions['u_0_prime']]
+        t_0 = self._get_parameter('t_0', theta)
+        u_0 = self._get_parameter('u_0', theta)
+        u_0_prime = self._get_parameter('u_0_prime', theta)
+
+        if u_0_prime is None:
+            return u_0 + (1 - torch.exp(-t + t_0)) * output_tensor
+        else:
+            return u_0 + (t - t_0) * u_0_prime + ((1 - torch.exp(-t + t_0)) ** 2) * output_tensor
 
-        if 't_0' in self.bundle_conditions:
 
-            t_0 = theta[self.bundle_conditions['t_0']]
+class BundleDirichletBVP(BaseCondition, _BundleConditionMixin):
+    r"""A double-ended Dirichlet boundary condition: :math:`u(t_0)=u_0` and :math:`u(t_1)=u_1`.
 
-            if self.u_0_prime is None and 'u_0_prime' not in self.bundle_conditions:
-                return u_0 + (t - t_0) * output_tensor
-            else:
-                return u_0 + (t - t_0) * u_0_prime + ((t - t_0) ** 2) * output_tensor
-        else:
-            if self.u_0_prime is None and 'u_0_prime' not in self.bundle_conditions:
-                return u_0 + (1 - torch.exp(-t + t_0)) * output_tensor
-            else:
-                return u_0 + (t - t_0) * u_0_prime + ((1 - torch.exp(-t + t_0)) ** 2) * output_tensor
+    :param t_0: The initial time. Ignored if 't_0' is present in bundle_param_lookup.
+    :type t_0: float
+    :param u_0: The initial value of :math:`u`. :math:`u(t_0)=u_0`. Ignored if 'u_0' is present in bundle_param_lookup.
+    :type u_0: float
+    :param t_1: The final time. Ignored if 't_1' is present in bundle_param_lookup.
+    :type t_1: float
+    :param u_1: The initial value of :math:`u`. :math:`u(t_1)=u_1`. Ignored if 'u_1' is present in bundle_param_lookup.
+    :type u_1: float
+    :param bundle_param_lookup: See _BundleConditionMixin for details. Allowed keys are 't_0', 'u_0', 't_1', and 'u_1'.
+    :type bundle_param_lookup: Dict[str, int]
+    """
+
+    @deprecated_alias(bundle_conditions='bundle_param_lookup')
+    def __init__(self, t_0, u_0, t_1, u_1, bundle_param_lookup=None):
+        BaseCondition.__init__(self)
+        _BundleConditionMixin.__init__(
+            self,
+            bundle_param_lookup=bundle_param_lookup,
+            allowed_params=['t_0', 'u_0', 't_1', 'u_1'],
+        )
+        self.t_0, self.u_0, self.t_1, self.u_1 = t_0, u_0, t_1, u_1
+
+    def parameterize(self, output_tensor, t, *theta):
+        r"""Re-parameterizes outputs such that the Dirichlet condition is satisfied on both ends of the domain.
+
+        The re-parameterization is
+        :math:`\\displaystyle u(t)=(1-\\tilde{t})u_0+\\tilde{t}u_1+\\left(1-e^{(1-\\tilde{t})\\tilde{t}}\\right)\\mathrm{ANN}(t)`,
+        where :math:`\\displaystyle \\tilde{t} = \\frac{t-t_0}{t_1-t_0}` and :math:`\\mathrm{ANN}` is the neural network.
+
+        :param output_tensor: Output of the neural network.
+        :type output_tensor: `torch.Tensor`
+        :param t: Input to the neural network; i.e., sampled time-points or another independent variable.
+        :type t: `torch.Tensor`
+        :param theta: Bundle parameters that potentially override default parameters.
+        :type theta: Tuple[`torch.Tensor`]
+        :return: The re-parameterized output of the network.
+        :rtype: `torch.Tensor`
+        """
+        u_0 = self._get_parameter('u_0', theta)
+        u_1 = self._get_parameter('u_1', theta)
+        t_0 = self._get_parameter('t_0', theta)
+        t_1 = self._get_parameter('t_1', theta)
+
+        t_tilde = (t - t_0) / (t_1 - t_0)
+        return u_0 * (1 - t_tilde) + u_1 * t_tilde + (1 - torch.exp((1 - t_tilde) * t_tilde)) * output_tensor
 
 
 class DirichletBVP(BaseCondition):
diff --git a/neurodiffeq/generators.py b/neurodiffeq/generators.py
index c2264cc..cdad4e8 100644
--- a/neurodiffeq/generators.py
+++ b/neurodiffeq/generators.py
@@ -804,7 +804,7 @@ class MeshGenerator(BaseGenerator):
     All sub-generators must return tensors of the same shape, or a tuple of tensors of the same shape.
     The number of tensors returned by each sub-generator can be different, but the intent behind
     this class is to create an N dimensional generator from several 1 dimensional generators, so each input generator
-    should represents one of the dimensions of your problem. An exception is made for
+    should represent one of the dimensions of your problem. An exception is made for
     using a ``MeshGenerator`` as one of the inputs of another ``MeshGenerator``. In that case the original
     meshed generators are extracted from the input ``MeshGenerator``, and the final mesh is used using those
     (e.g ``MeshGenerator(MeshGenerator(g1, g2), g3)`` is equivalent to ``MeshGenerator(g1, g2, g3)``, where
diff --git a/neurodiffeq/solvers.py b/neurodiffeq/solvers.py
index 13f03f6..c97918e 100644
--- a/neurodiffeq/solvers.py
+++ b/neurodiffeq/solvers.py
@@ -1196,18 +1196,25 @@ class BundleSolver1D(BaseSolver):
     :type t_max: float, optional
     :param theta_min:
         Lower bound of input (parameters and/or conditions). If conditions are included in the bundle,
-        the order should match the one inferred by the values of the ``bundle_conditions`` input
+        the order should match the one inferred by the values of the ``bundle_param_lookup`` input
         in the ``neurodiffeq.conditions.BundleIVP``.
         Defaults to None.
         Ignored if ``train_generator`` and ``valid_generator`` are both set.
     :type theta_min: float or tuple, optional
     :param theta_max:
         Upper bound of input (parameters and/or conditions). If conditions are included in the bundle,
-        the order should match the one inferred by the values of the ``bundle_conditions`` input
+        the order should match the one inferred by the values of the ``bundle_param_lookup`` input
         in the ``neurodiffeq.conditions.BundleIVP``.
         Defaults to None.
         Ignored if ``train_generator`` and ``valid_generator`` are both set.
     :type theta_max: float or tuple, optional
+    :param eq_param_index:
+        Index (or indices) of bundle parameter that appears in the equation.
+        E.g., if there are 5 bundle parameters generated and the first (index 0) and last (index 4) parameters are used
+        in the equation, then eq_param_index should be (0, 4).
+        The signature of the original equation should have, in addition to functions to solve for and coordinates,
+        2 more parameters, corresponding to bundle parameters indexed at 0 and 4, in that order.
+    :type eq_param_index: int or tuple[int], optional
     :param nets:
         List of neural networks for parameterized solution.
         If provided, length of ``nets`` must equal that of ``conditions``
@@ -1216,11 +1223,13 @@ class BundleSolver1D(BaseSolver):
         Generator for sampling training points,
         which must provide a ``.get_examples()`` method and a ``.size`` field.
         ``train_generator`` must be specified if ``t_min`` and ``t_max`` are not set.
+        If provided, the generator must generate bundle parameters too.
     :type train_generator: `neurodiffeq.generators.BaseGenerator`, optional
     :param valid_generator:
         Generator for sampling validation points,
         which must provide a ``.get_examples()`` method and a ``.size`` field.
         ``valid_generator`` must be specified if ``t_min`` and ``t_max`` are not set.
+        If provided, the generator must generate bundle parameters too.
     :type valid_generator: `neurodiffeq.generators.BaseGenerator`, optional
     :param analytic_solutions:
         Analytical solutions to be compared with neural net solutions.
@@ -1278,7 +1287,7 @@ class BundleSolver1D(BaseSolver):
     """
 
     def __init__(self, ode_system, conditions, t_min, t_max,
-                 theta_min=None, theta_max=None,
+                 theta_min=None, theta_max=None, eq_param_index=(),
                  nets=None, train_generator=None, valid_generator=None, analytic_solutions=None, optimizer=None,
                  loss_fn=None, n_batches_train=1, n_batches_valid=4, metrics=None, n_output_units=1,
                  # deprecated arguments are listed below
@@ -1292,16 +1301,21 @@ class BundleSolver1D(BaseSolver):
 
         if isinstance(theta_min, (float, int)):
             theta_min = (theta_min,)
+        elif theta_min is None:
+            theta_min = ()
 
         if isinstance(theta_max, (float, int)):
             theta_max = (theta_max,)
+        elif theta_max is None:
+            theta_max = ()
 
-        if theta_min is None and theta_max is None:
-            r_min = (t_min,)
-            r_max = (t_max,)
-        else:
-            r_min = (t_min,) + theta_min
-            r_max = (t_max,) + theta_max
+        if len(theta_min) != len(theta_max):
+            raise ValueError(
+                f"length of theta_min and theta_max must be equal, " f"got {len(theta_min)} != {len(theta_max)}"
+            )
+
+        r_min = (t_min,) + tuple(theta_min)
+        r_max = (t_max,) + tuple(theta_max)
 
         n_input_units = len(r_min)
 
@@ -1315,26 +1329,28 @@ class BundleSolver1D(BaseSolver):
             for i in range(n_input_units - 1):
                 valid_generator ^= Generator1D(32, t_min=r_min[i + 1], t_max=r_max[i + 1], method='equally-spaced')
 
-        non_var = []
-        for c in conditions:
-            try:
-                bc = tuple(c.bundle_conditions.values())
-            except Exception:
-                bc = ()
-            for index in bc:
-                non_var.append(index)
-        non_var = list(set(non_var))
-
-        def non_var_filter(*variables):
-            variables = list(variables)
-            for i, n in enumerate(non_var):
-                del variables[len(conditions) + 1 + n - i]
-            return ode_system(*variables)
-
         self.r_min, self.r_max = r_min, r_max
 
+        # number of functions equals number of conditions supplied
+        N_FUNCTIONS = len(conditions)
+        # there is only 1 coordinate (usually time `t`) for ODEs
+        N_COORDS = 1
+
+        # Note: It is intentionally design in this way where `eq_param_index` and `self.eq_param_index`
+        # both contain values offset by `N_FUNCTIONS + N_COORDS`. The first one (not bound to `self`) is used for the
+        # `eq_param_filter` closure, while the second one (bound to `self`) is used for `_get_internal_variables()`.
+        eq_param_index = tuple(N_FUNCTIONS + N_COORDS + idx for idx in eq_param_index)
+        self.eq_param_index = eq_param_index
+
+        # TODO: check and warn if there are variables neither used by conditions nor by equations
+
+        def _diff_eqs_wrapper(*variables):
+            funcs_and_coords = variables[:N_FUNCTIONS + N_COORDS]
+            eq_params = tuple(variables[idx] for idx in eq_param_index)
+            return ode_system(*funcs_and_coords, *eq_params)
+
         super(BundleSolver1D, self).__init__(
-            diff_eqs=non_var_filter,
+            diff_eqs=_diff_eqs_wrapper,
             conditions=conditions,
             nets=nets,
             train_generator=train_generator,
@@ -1387,6 +1403,7 @@ class BundleSolver1D(BaseSolver):
         available_variables.update({
             'r_min': self.r_min,
             'r_max': self.r_max,
+            'eq_param_index': self.eq_param_index,
         })
         return available_variables
 
diff --git a/tests/test_conditions.py b/tests/test_conditions.py
index aa7b58e..6f4e5a8 100644
--- a/tests/test_conditions.py
+++ b/tests/test_conditions.py
@@ -8,6 +8,7 @@ from neurodiffeq.conditions import IVP
 from neurodiffeq.conditions import BundleIVP
 from neurodiffeq.conditions import EnsembleCondition
 from neurodiffeq.conditions import DirichletBVP
+from neurodiffeq.conditions import BundleDirichletBVP
 from neurodiffeq.conditions import DirichletBVP2D
 from neurodiffeq.conditions import DirichletBVPSpherical
 from neurodiffeq.conditions import InfDirichletBVPSpherical
@@ -164,7 +165,8 @@ def test_ivp(x0, y0, y1, ones, net11):
     assert all_close(diff(y, x), y1), "y'(x_0) != y'_0"
 
 
-def test_bundleivp(x0, y0, y1, ones, lin, net11, net21, net31, net41):
+# TODO parameterize the following test cases
+def test_bundle_ivp(x0, y0, y1, ones, lin, net11, net21, net31, net41):
     # Regular IVP with no bundle:
     x = x0 * ones
     cond = BundleIVP(x0, y0)
@@ -178,24 +180,24 @@ def test_bundleivp(x0, y0, y1, ones, lin, net11, net21, net31, net41):
 
     # Bundle in u_0:
     y_bundle = y0 * lin
-    cond = BundleIVP(t_0=x0, bundle_conditions={'u_0': 0})
+    cond = BundleIVP(t_0=x0, bundle_param_lookup={'u_0': 0})
     y = cond.enforce(net21, x, y_bundle)
     assert torch.isclose(y, y0 * lin).all(), "y(x_0) != y_0"
 
-    cond = BundleIVP(t_0=x0, u_0_prime=y1, bundle_conditions={'u_0': 0})
+    cond = BundleIVP(t_0=x0, u_0_prime=y1, bundle_param_lookup={'u_0': 0})
     y = cond.enforce(net21, x, y_bundle)
     assert torch.isclose(y, y0 * lin).all(), "y(x_0) != y_0"
     assert all_close(diff(y, x), y1), "y'(x_0) != y'_0"
 
     # Bundle in u_0_prime:
     y_prime_bundle = y1 * lin
-    cond = BundleIVP(t_0=x0, u_0=y0, bundle_conditions={'u_0_prime': 0})
+    cond = BundleIVP(t_0=x0, u_0=y0, bundle_param_lookup={'u_0_prime': 0})
     y = cond.enforce(net21, x, y_prime_bundle)
     assert all_close(y, y0), "y(x_0) != y_0"
     assert torch.isclose(diff(y, x), y1 * lin).all(), "y'(x_0) != y'_0"
 
     # Bundle in u_0 and u_0_prime:
-    cond = BundleIVP(t_0=x0, bundle_conditions={'u_0': 0, 'u_0_prime': 1})
+    cond = BundleIVP(t_0=x0, bundle_param_lookup={'u_0': 0, 'u_0_prime': 1})
     y = cond.enforce(net31, x, y_bundle, y_prime_bundle)
     assert torch.isclose(y, y0 * lin).all(), "y(x_0) != y_0"
     assert torch.isclose(diff(y, x), y1 * lin).all(), "y'(x_0) != y'_0"
@@ -203,32 +205,38 @@ def test_bundleivp(x0, y0, y1, ones, lin, net11, net21, net31, net41):
     # Bundle in t_0:
     x = x0 * lin
     x_bundle = x0 * lin
-    cond = BundleIVP(u_0=y0, bundle_conditions={'t_0': 0})
+    cond = BundleIVP(u_0=y0, bundle_param_lookup={'t_0': 0})
     y = cond.enforce(net21, x, x_bundle)
     assert torch.isclose(y, y0 * ones).all(), "y(x_0) != y_0"
 
-    cond = BundleIVP(u_0=y0, u_0_prime=y1, bundle_conditions={'t_0': 0})
+    cond = BundleIVP(u_0=y0, u_0_prime=y1, bundle_param_lookup={'t_0': 0})
     y = cond.enforce(net21, x, x_bundle)
     assert all_close(y, y0), "y(x_0) != y_0"
     assert all_close(diff(y, x), y1), "y'(x_0) != y'_0"
 
     # Bundle in t_0 and u_0:
-    cond = BundleIVP(bundle_conditions={'t_0': 0, 'u_0': 1})
+    cond = BundleIVP(bundle_param_lookup={'t_0': 0, 'u_0': 1})
     y = cond.enforce(net31, x, x_bundle, y_bundle)
     assert torch.isclose(y, y0 * lin).all(), "y(x_0) != y_0"
 
-    cond = BundleIVP(u_0_prime=y1, bundle_conditions={'t_0': 0, 'u_0': 1})
+    cond = BundleIVP(u_0_prime=y1, bundle_param_lookup={'t_0': 0, 'u_0': 1})
     y = cond.enforce(net31, x, x_bundle, y_bundle)
     assert torch.isclose(y, y0 * lin).all(), "y(x_0) != y_0"
     assert all_close(diff(y, x), y1), "y'(x_0) != y'_0"
 
     # Bundle in t_0, u_0 and u_0_prime:
-    cond = BundleIVP(bundle_conditions={'t_0': 0, 'u_0': 1, 'u_0_prime': 2})
+    cond = BundleIVP(bundle_param_lookup={'t_0': 0, 'u_0': 1, 'u_0_prime': 2})
     y = cond.enforce(net41, x, x_bundle, y_bundle, y_prime_bundle)
     assert torch.isclose(y, y0 * lin).all(), "y(x_0) != y_0"
     assert torch.isclose(diff(y, x), y1 * lin).all(), "y'(x_0) != y'_0"
 
 
+@pytest.mark.parametrize(argnames='illegal_name', argvalues=['magic', 't0', 'u1', 'u1prime'])
+def test_bundle_ivp__disallowed_params(x0, y0, y1, illegal_name):
+    with pytest.raises(ValueError):
+        BundleIVP(x0, y0, y1, bundle_param_lookup={illegal_name: 0})
+
+
 def test_ivp_legacy_signature():
     with warns(FutureWarning):
         IVP(0, x_0=1)
@@ -279,7 +287,50 @@ def test_dirichlet_bvp(x0, x1, y0, y1, ones, net11):
     assert all_close(y, y1), "y(x_1) != y_1"
 
 
-def test_bvp_legacy_signature():
+@pytest.mark.parametrize(argnames='t_0_bundle', argvalues=[True, False])
+@pytest.mark.parametrize(argnames='u_0_bundle', argvalues=[True, False])
+@pytest.mark.parametrize(argnames='t_1_bundle', argvalues=[True, False])
+@pytest.mark.parametrize(argnames='u_1_bundle', argvalues=[True, False])
+def test_bundle_dirichlet_bvp(x0, y0, x1, y1, ones, t_0_bundle, u_0_bundle, t_1_bundle, u_1_bundle):
+    __params = ['t_0', 'u_0', 't_1', 'u_1']
+    __local_vars = locals()
+    bundled_params = [p for p in __params if __local_vars[p + '_bundle']]
+    bundle_vals = [torch.rand(ones.shape) for _ in bundled_params]
+    bundle_param_lookup = {p: idx for idx, p in enumerate(bundled_params)}
+
+    cond = BundleDirichletBVP(x0, y0, x1, y1, bundle_param_lookup=bundle_param_lookup)
+
+    n_bundle_parameters = sum([t_0_bundle, u_0_bundle, t_1_bundle, u_1_bundle])
+    net = FCNN(n_input_units=1 + n_bundle_parameters, n_output_units=1)
+
+    if t_0_bundle:
+        y = cond.enforce(net, bundle_vals[bundle_param_lookup['t_0']], *bundle_vals)
+    else:
+        y = cond.enforce(net, x0 * ones, *bundle_vals)
+    z = bundle_vals[bundle_param_lookup['u_0']] if 'u_0' in bundle_param_lookup else y0
+    assert all_close(y, z), 'y(x_0) != y_0'
+
+    if t_1_bundle:
+        y = cond.enforce(net, bundle_vals[bundle_param_lookup['t_1']], *bundle_vals)
+    else:
+        y = cond.enforce(net, x1 * ones, *bundle_vals)
+    z = bundle_vals[bundle_param_lookup['u_1']] if 'u_1' in bundle_param_lookup else y1
+    assert all_close(y, z), 'y(x_1) != y_1'
+
+
+@pytest.mark.parametrize(argnames='illegal_name', argvalues=['magic', 'u0', 't0', 'u1', 't1'])
+def test_bundle_dirichlet_bvp__disallowed_params(x0, y0, x1, y1, illegal_name):
+    with pytest.raises(ValueError):
+        BundleDirichletBVP(x0, y0, x1, y1, bundle_param_lookup={illegal_name: 0})
+
+
+def test_bundle_conditions__legacy_signature(x0, y0, x1, y1):
+    with pytest.warns(FutureWarning):
+        BundleIVP(x0, y0, bundle_conditions={'t_0': 0})
+    with pytest.warns(FutureWarning):
+        BundleDirichletBVP(x0, y0, x1, y1, bundle_conditions={'t_0': 0})
+
+def test_bvp__legacy_signature():
     with warns(FutureWarning):
         DirichletBVP(t_0=0, t_1=0, x_0=0, x_1=0)
     with warns(FutureWarning):