CUQIpy 1.3.0.post0.dev298__tar.gz → 1.3.0.post0.dev362__tar.gz

{cuqipy-1.3.0.post0.dev298 → cuqipy-1.3.0.post0.dev362}/CUQIpy.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: CUQIpy
-Version: 1.3.0.post0.dev298
+Version: 1.3.0.post0.dev362
 Summary: Computational Uncertainty Quantification for Inverse problems in Python
 Maintainer-email: "Nicolai A. B. Riis" <nabr@dtu.dk>, "Jakob S. Jørgensen" <jakj@dtu.dk>, "Amal M. Alghamdi" <amaal@dtu.dk>, Chao Zhang <chaz@dtu.dk>
 License:                                  Apache License

{cuqipy-1.3.0.post0.dev298 → cuqipy-1.3.0.post0.dev362}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: CUQIpy
-Version: 1.3.0.post0.dev298
+Version: 1.3.0.post0.dev362
 Summary: Computational Uncertainty Quantification for Inverse problems in Python
 Maintainer-email: "Nicolai A. B. Riis" <nabr@dtu.dk>, "Jakob S. Jørgensen" <jakj@dtu.dk>, "Amal M. Alghamdi" <amaal@dtu.dk>, Chao Zhang <chaz@dtu.dk>
 License:                                  Apache License

{cuqipy-1.3.0.post0.dev298 → cuqipy-1.3.0.post0.dev362}/cuqi/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-09-04T09:48:37+0300",
+ "date": "2025-09-12T10:23:05+0300",
  "dirty": false,
  "error": null,
- "full-revisionid": "a0d227d0981362a09d24dc976ad65459ba8bf798",
- "version": "1.3.0.post0.dev298"
+ "full-revisionid": "f6a73b0b32186614fe5451781567c1abcd48452a",
+ "version": "1.3.0.post0.dev362"
 }
 '''  # END VERSION_JSON
 

{cuqipy-1.3.0.post0.dev298 → cuqipy-1.3.0.post0.dev362}/cuqi/distribution/_distribution.py

@@ -105,7 +105,7 @@ class Distribution(Density, ABC):
                 f"Inconsistent distribution geometry attribute {self._geometry} and inferred "
                 f"dimension from distribution variables {inferred_dim}."
             )
-        
+
         # If Geometry dimension is None, update it with the inferred dimension
         if inferred_dim and self._geometry.par_dim is None: 
             self.geometry = inferred_dim
@@ -117,7 +117,7 @@ class Distribution(Density, ABC):
         # We do not use self.name to potentially infer it from python stack.
         if self._name: 
             self._geometry._variable_name = self._name
-            
+
         return self._geometry
 
     @geometry.setter
@@ -160,7 +160,7 @@ class Distribution(Density, ABC):
                     f"{self.logd.__qualname__}: To evaluate the log density all conditioning variables and main"
                     f" parameter must be specified. Conditioning variables are: {cond_vars}"
                 )
-            
+
             # Check if all conditioning variables are specified
             all_cond_vars_specified = all([key in kwargs for key in cond_vars])
             if not all_cond_vars_specified:
@@ -168,7 +168,7 @@ class Distribution(Density, ABC):
                     f"{self.logd.__qualname__}: To evaluate the log density all conditioning variables must be"
                     f" specified. Conditioning variables are: {cond_vars}"
                 )
-            
+
             # Extract exactly the conditioning variables from kwargs
             cond_kwargs = {key: kwargs[key] for key in cond_vars}
 
@@ -186,7 +186,7 @@ class Distribution(Density, ABC):
         # Not conditional distribution, simply evaluate log density directly
         else:
             return super().logd(*args, **kwargs)
-        
+
     def _logd(self, *args):
         return self.logpdf(*args) # Currently all distributions implement logpdf so we simply call this method.
 
@@ -216,7 +216,7 @@ class Distribution(Density, ABC):
         # Get samples from the distribution sample method
         s = self._sample(N,*args,**kwargs)
 
-        #Store samples in cuqi samples object if more than 1 sample
+        # Store samples in cuqi samples object if more than 1 sample
         if N==1:
             if len(s) == 1 and isinstance(s,np.ndarray): #Extract single value from numpy array
                 s = s.ravel()[0]
@@ -264,7 +264,7 @@ class Distribution(Density, ABC):
         # Go through every mutable variable and assign value from kwargs if present
         for var_key in mutable_vars:
 
-            #If keyword directly specifies new value of variable we simply reassign
+            # If keyword directly specifies new value of variable we simply reassign
             if var_key in kwargs:
                 setattr(new_dist, var_key, kwargs.get(var_key))
                 processed_kwargs.add(var_key)
@@ -291,9 +291,18 @@ class Distribution(Density, ABC):
 
                 elif len(var_args)>0:                      #Some keywords found
                     # Define new partial function with partially defined args
-                    func = partial(var_val, **var_args)
+                    if (
+                        hasattr(var_val, "_supports_partial_eval")
+                        and var_val._supports_partial_eval
+                    ):
+                        func = var_val(**var_args)
+                    else:
+                        # If the callable does not support partial evaluation,
+                        # we use the partial function to set the variable
+                        func = partial(var_val, **var_args)
+
                     setattr(new_dist, var_key, func)
-                
+
                 # Store processed keywords
                 processed_kwargs.update(var_args.keys())
 
@@ -329,7 +338,7 @@ class Distribution(Density, ABC):
 
     def get_conditioning_variables(self):
         """Return the conditioning variables of this distribution (if any)."""
-        
+
         # Get all mutable variables
         mutable_vars = self.get_mutable_variables()
 
@@ -338,7 +347,7 @@ class Distribution(Density, ABC):
 
         # Add any variables defined through callable functions
         cond_vars += get_indirect_variables(self)
-        
+
         return cond_vars
 
     def get_mutable_variables(self):
@@ -347,10 +356,10 @@ class Distribution(Density, ABC):
         # If mutable variables are already cached, return them
         if hasattr(self, '_mutable_vars'):
             return self._mutable_vars
-        
+
         # Define list of ignored attributes and properties
         ignore_vars = ['name', 'is_symmetric', 'geometry', 'dim']
-        
+
         # Get public attributes
         attributes = get_writeable_attributes(self)
 
@@ -396,7 +405,7 @@ class Distribution(Density, ABC):
                         raise ValueError(f"{self._condition.__qualname__}: {ordered_keys[index]} passed as both argument and keyword argument.\nArguments follow the listed conditioning variable order: {self.get_conditioning_variables()}")
                     kwargs[ordered_keys[index]] = arg
         return kwargs
-    
+
     def _check_geometry_consistency(self):
         """ Checks that the geometry of the distribution is consistent by calling the geometry property. Should be called at the end of __init__ of subclasses. """
         self.geometry
@@ -411,4 +420,4 @@ class Distribution(Density, ABC):
     def rv(self):
         """ Return a random variable object representing the distribution. """
         from cuqi.experimental.algebra import RandomVariable
-        return RandomVariable(self)
+        return RandomVariable(self)

{cuqipy-1.3.0.post0.dev298 → cuqipy-1.3.0.post0.dev362}/cuqi/likelihood/_likelihood.py

@@ -212,4 +212,4 @@ class UserDefinedLikelihood(object):
         return get_non_default_args(self.logpdf_func)
 
     def __repr__(self) -> str:
-        return "CUQI {} function. Parameters {}.".format(self.__class__.__name__,self.get_parameter_names())
+        return "CUQI {} function. Parameters {}.".format(self.__class__.__name__,self.get_parameter_names())

{cuqipy-1.3.0.post0.dev298 → cuqipy-1.3.0.post0.dev362}/cuqi/model/_model.py

@@ -132,6 +132,10 @@ class Model(object):
         print(model(1, 1))
         print(model.gradient(np.array([1]), 1, 1))
     """
+
+    _supports_partial_eval = True
+    """Flag indicating that partial evaluation of Model objects is supported, i.e., calling the model object with only some of the inputs specified returns a model that can be called with the remaining inputs."""
+
     def __init__(self, forward, range_geometry, domain_geometry, gradient=None, jacobian=None):
 
         # Check if input is callable
@@ -311,7 +315,12 @@ class Model(object):
                 "Gradient needs to be callable function or tuple of callable functions."
             )
 
-        expected_func_non_default_args = self._non_default_args
+        expected_func_non_default_args = (
+            self._non_default_args
+            if not hasattr(self, "_original_non_default_args")
+            else self._original_non_default_args
+        )
+
         if func_type.lower() == "gradient":
             # prepend 'direction' to the expected gradient non default args
             expected_func_non_default_args = [
@@ -613,52 +622,43 @@ class Model(object):
         if non_default_args is None:
             non_default_args = self._non_default_args
 
-        # If any args are given, add them to kwargs
+        # Either args or kwargs can be provided but not both
+        if len(args) > 0 and len(kwargs) > 0:
+            raise ValueError(
+                "The "
+                + map_name.lower()
+                + " input is specified both as positional and keyword arguments. This is not supported."
+            )
+
+        len_input = len(args) + len(kwargs)
+
+        # If partial evaluation, make sure input is not of type Samples
+        if len_input < len(non_default_args):
+            # If the argument is a Sample object, splitting or partial
+            # evaluation of the model is not supported
+            temp_args = args if len(args) > 0 else list(kwargs.values())
+            if any(isinstance(arg, Samples) for arg in temp_args):
+                raise ValueError(("When using Samples objects as input, the"
+                                +" user should provide a Samples object for"
+                                +f" each non_default_args {non_default_args}"
+                                +" of the model. That is, partial evaluation"
+                                +" or splitting is not supported for input"
+                                +" of type Samples."))
+
+        # If args are given, add them to kwargs
         if len(args) > 0:
-            if len(kwargs) > 0:
-                raise ValueError(
-                    "The "
-                    + map_name.lower()
-                    + " input is specified both as positional and keyword arguments. This is not supported."
-                )
 
-            appending_error_message = ""
             # Check if the input is for multiple input case and is stacked,
             # then split it
-            if len(args)==1 and len(non_default_args)>1:
-                # If the argument is a Sample object, splitting is not supported
-                if isinstance(args[0], Samples):
-                    raise ValueError(
-                        "The "
-                        + map_name.lower()
-                        + f" input is specified by a Samples object that cannot be split into multiple arguments corresponding to the non_default_args {non_default_args}."
-                    )
-                split_succeeded, split_args = self._is_stacked_args(*args, is_par=is_par)
-                if split_succeeded:
-                    args = split_args
-                else:
-                    appending_error_message = (
-                        " Additionally, the "
-                        + map_name.lower()
-                        + f" input is specified by a single argument that cannot be split into multiple arguments matching the expected non_default_args {non_default_args}."
-                    )
-
-            # Check if the number of args does not match the number of
-            # non_default_args of the model
-            if len(args) != len(non_default_args):
-                raise ValueError(
-                    "The number of positional arguments does not match the number of non-default arguments of the "
-                    + map_name.lower()
-                    + "."
-                    + appending_error_message
-                )
+            if len(args) < len(non_default_args):
+                args = self._split_in_case_of_stacked_args(*args, is_par=is_par)
 
             # Add args to kwargs following the order of non_default_args
             for idx, arg in enumerate(args):
                 kwargs[non_default_args[idx]] = arg
-
+    
         # Check kwargs matches non_default_args
-        if set(list(kwargs.keys())) != set(non_default_args):
+        if not (set(list(kwargs.keys())) <= set(non_default_args)):
             if map_name == "gradient":
                 error_msg = f"The gradient input is specified by a direction and keywords arguments {list(kwargs.keys())} that does not match the non_default_args of the model {non_default_args}."
             else:
@@ -673,53 +673,41 @@ class Model(object):
             raise ValueError(error_msg)
 
         # Make sure order of kwargs is the same as non_default_args
-        kwargs = {k: kwargs[k] for k in non_default_args}
+        kwargs = {k: kwargs[k] for k in non_default_args if k in kwargs}
 
         return kwargs
 
-    def _is_stacked_args(self, *args, is_par=True):
-        """Private function that checks if the input arguments are stacked
-        and splits them if they are."""
-        # Length of args should be 1 if the input is stacked (no partial
-        # stacking is supported)
-        if len(args) > 1:
-            return False, args
-
-        # Type of args should be parameter
-        if not is_par:
-            return False, args
+    def _split_in_case_of_stacked_args(self, *args, is_par=True):
+        """Private function that checks if the input args is a stacked
+        CUQIarray or numpy array and splits it into multiple arguments based on
+        the domain geometry of the model. Otherwise, it returns the input args
+        unchanged."""
 
-        # args[0] should be numpy array or CUQIarray
+        # Check conditions for splitting and split if all conditions are met
         is_CUQIarray = isinstance(args[0], CUQIarray)
         is_numpy_array = isinstance(args[0], np.ndarray)
-        if not is_CUQIarray and not is_numpy_array:
-            return False, args
 
-        # Shape of args[0] should be (domain_dim,)
-        if not args[0].shape == (self.domain_dim,):
-            return False, args
-
-        # Ensure domain geometry is _ProductGeometry
-        if not isinstance(
-            self.domain_geometry, cuqi.experimental.geometry._ProductGeometry
-        ):
-            return False, args
-
-        # Split the stacked input
-        split_args = np.split(args[0], self.domain_geometry.stacked_par_split_indices)
-
-        # Covert split args to CUQIarray if input is CUQIarray
-        if is_CUQIarray:
-            split_args = [
-                CUQIarray(arg, is_par=True, geometry=self.domain_geometry.geometries[i])
-                for i, arg in enumerate(split_args)
-            ]
+        if ((is_CUQIarray or is_numpy_array) and
+           is_par and
+           len(args) == 1 and
+           args[0].shape == (self.domain_dim,) and
+           isinstance(self.domain_geometry, cuqi.experimental.geometry._ProductGeometry)):
+            # Split the stacked input
+            split_args = np.split(args[0], self.domain_geometry.stacked_par_split_indices)
+            # Convert split args to CUQIarray if input is CUQIarray
+            if is_CUQIarray:
+                split_args = [
+                    CUQIarray(arg, is_par=True, geometry=self.domain_geometry.geometries[i])
+                    for i, arg in enumerate(split_args)
+                ]
+            return split_args
 
-        return True, split_args
+        else:
+            return args
 
     def forward(self, *args, is_par=True, **kwargs):
         """ Forward function of the model.
-        
+
         Forward converts the input to function values (if needed) using the domain geometry of the model. Then it applies the forward operator to the function values and converts the output to parameters using the range geometry of the model.
 
         Parameters
@@ -733,7 +721,7 @@ class Model(object):
             If True, the inputs in `args` or `kwargs` are assumed to be parameters.
             If False, the inputs in `args` or `kwargs` are assumed to be function values.
             If `is_par` is a tuple of bools, the inputs are assumed to be parameters or function values based on the corresponding boolean value in the tuple.
-        
+
         **kwargs : keyword arguments
             keyword arguments for the forward operator. The forward operator input can be specified as either positional arguments or keyword arguments but not both.
 
@@ -750,19 +738,31 @@ class Model(object):
         kwargs = self._parse_args_add_to_kwargs(
             *args, **kwargs, is_par=is_par, map_name="model"
         )
-
-        # extract args from kwargs
+        # Extract args from kwargs
         args = list(kwargs.values())
 
+        if len(kwargs) == 0:
+            return self
+
+        partial_arguments = len(kwargs) < len(self._non_default_args)
+
         # If input is a distribution, we simply change the parameter name of
         # model to match the distribution name
         if all(isinstance(x, cuqi.distribution.Distribution)
                for x in kwargs.values()):
+            if partial_arguments:
+                raise ValueError(
+                    "Partial evaluation of the model is not supported for distributions."
+                )
             return self._handle_case_when_model_input_is_distributions(kwargs)
 
         # If input is a random variable, we handle it separately
         elif all(isinstance(x, cuqi.experimental.algebra.RandomVariable)
                for x in kwargs.values()):
+            if partial_arguments:
+                raise ValueError(
+                    "Partial evaluation of the model is not supported for random variables."
+                )
             return self._handle_case_when_model_input_is_random_variables(kwargs)
 
         # If input is a Node from internal abstract syntax tree, we let the Node handle the operation
@@ -772,6 +772,21 @@ class Model(object):
         elif any(isinstance(args_i, cuqi.experimental.algebra.Node) for args_i in args):
             return NotImplemented
 
+        # if input is partial, we create a new model with the partial input
+        if partial_arguments:
+            # Create is_par_partial from the is_par to contain only the relevant parts
+            if isinstance(is_par, (list, tuple)):
+                is_par_partial = tuple(
+                    is_par[i]
+                    for i in range(self.number_of_inputs)
+                    if self._non_default_args[i] in kwargs.keys()
+                )
+            else:
+                is_par_partial = is_par
+            # Build a partial model with the given kwargs
+            partial_model = self._build_partial_model(kwargs, is_par_partial)
+            return partial_model
+
         # Else we apply the forward operator
         # if model has _original_non_default_args, we use it to replace the
         # kwargs keys so that it matches self._forward_func signature
@@ -797,6 +812,126 @@ class Model(object):
         else:
             return False
 
+    def _build_partial_model(self, kwargs, is_par):
+        """Private function that builds a partial model substituting the given
+        keyword arguments with their values. The created partial model will have
+        as inputs the non-default arguments that are not in the kwargs."""
+
+        # Extract args from kwargs
+        args = list(kwargs.values())
+
+        # Define original_non_default_args which represents the complete list of
+        # non-default arguments of the forward function.
+        original_non_default_args = (
+            self._original_non_default_args
+            if hasattr(self, "_original_non_default_args")
+            else self._non_default_args
+        )
+
+        if hasattr(self, "_original_non_default_args"):
+            # Split the _original_non_default_args into two lists:
+            # 1. reduced_original_non_default_args: the _original_non_default_args
+            # corresponding to the _non_default_args that are not in kwargs
+            # 2. substituted_non_default_args: the _original_non_default_args
+            # corresponding to the _non_default_args that are in kwargs
+            reduced_original_non_default_args = [
+                original_non_default_args[i]
+                for i in range(self.number_of_inputs)
+                if self._non_default_args[i] not in kwargs.keys()
+            ]
+            substituted_non_default_args = [
+                original_non_default_args[i]
+                for i in range(self.number_of_inputs)
+                if self._non_default_args[i] in kwargs.keys()
+            ]
+            # Replace the keys in kwargs with the substituted_non_default_args
+            # so that the kwargs match the signature of the _forward_func
+            kwargs = {k: v for k, v in zip(substituted_non_default_args, args)}
+
+        # Create a partial domain geometry with the geometries corresponding
+        # to the non-default arguments that are not in kwargs (remaining
+        # unspecified inputs)
+        partial_domain_geometry = cuqi.experimental.geometry._ProductGeometry(
+            *[
+                self.domain_geometry.geometries[i]
+                for i in range(self.number_of_inputs)
+                if original_non_default_args[i] not in kwargs.keys()
+            ]
+        )
+
+        if len(partial_domain_geometry.geometries) == 1:
+            partial_domain_geometry = partial_domain_geometry.geometries[0]
+
+        # Create a domain geometry with the geometries corresponding to the
+        # non-default arguments that are specified
+        substituted_domain_geometry = cuqi.experimental.geometry._ProductGeometry(
+            *[
+                self.domain_geometry.geometries[i]
+                for i in range(self.number_of_inputs)
+                if original_non_default_args[i] in kwargs.keys()
+            ]
+        )
+
+        if len(substituted_domain_geometry.geometries) == 1:
+            substituted_domain_geometry = substituted_domain_geometry.geometries[0]
+
+        # Create new model with partial input
+        # First, we convert the input to function values
+        kwargs = self._2fun(geometry=substituted_domain_geometry, is_par=is_par, **kwargs)
+
+        # Second, we create a partial function for the forward operator
+        partial_forward = partial(self._forward_func, **kwargs)
+
+        # Third, if applicable, we create a partial function for the gradient
+        if isinstance(self._gradient_func, tuple):
+            # If gradient is a tuple, we create a partial function for each
+            # gradient function in the tuple
+            partial_gradient = tuple(
+                (
+                    partial(self._gradient_func[i], **kwargs)
+                    if self._gradient_func[i] is not None
+                    else None
+                )
+                for i in range(self.number_of_inputs)
+                if original_non_default_args[i] not in kwargs.keys()
+            )
+            if len(partial_gradient) == 1:
+                partial_gradient = partial_gradient[0]
+
+        elif callable(self._gradient_func):
+            raise NotImplementedError(
+                "Partial forward model is only supported for gradient/jacobian functions that are tuples of callable functions."
+            )
+
+        else:
+            partial_gradient = None
+
+        # Lastly, we create the partial model with the partial forward
+        # operator (we set the gradient function later)
+        partial_model = Model(
+            forward=partial_forward,
+            range_geometry=self.range_geometry,
+            domain_geometry=partial_domain_geometry,
+        )
+
+        # Set the _original_non_default_args (if applicable) and
+        # _stored_non_default_args of the partial model
+        if hasattr(self, "_original_non_default_args"):
+            partial_model._original_non_default_args = reduced_original_non_default_args
+        partial_model._stored_non_default_args = [
+            self._non_default_args[i]
+            for i in range(self.number_of_inputs)
+            if original_non_default_args[i] not in kwargs.keys()
+        ]
+
+        # Set the gradient function of the partial model
+        partial_model._check_correct_gradient_jacobian_form(
+            partial_gradient, "gradient"
+        )
+        partial_model._gradient_func = partial_gradient
+
+        return partial_model
+
     def _handle_case_when_model_input_is_distributions(self, kwargs):
         """Private function that handles the case of the input being a
         distribution or multiple distributions."""

{cuqipy-1.3.0.post0.dev298 → cuqipy-1.3.0.post0.dev362}/tests/test_joint_distribution.py

@@ -1,6 +1,7 @@
 import cuqi
 import numpy as np
 import pytest
+from .test_model import MultipleInputTestModel
 
 def test_joint_dist_dim_geometry():
     """ Test the dimension and geometry properties of a joint distribution """
@@ -393,3 +394,94 @@ def test_logd_consistency_when_conditioning(joint, variables):
 
         # Add current variable to the variables that need to be conditioned
         cond_vars[key] = value
+
+def test_joint_distribution_with_multiple_inputs_model_has_correct_parameter_names():
+    """Test that the joint distribution based on model with multiple inputs has
+    correct parameter names."""
+
+    test_model = MultipleInputTestModel.helper_build_three_input_test_model()
+    model = cuqi.model.Model(
+        test_model.forward_map,
+        gradient=test_model.gradient_form2,
+        domain_geometry=test_model.domain_geometry,
+        range_geometry=test_model.range_geometry,
+    )
+
+    # Create priors
+    x_dist = cuqi.distribution.Gaussian(
+        mean=np.zeros(3),
+        cov=np.eye(3))
+    y_dist = cuqi.distribution.Gaussian(
+        mean=np.zeros(2),
+        cov=np.eye(2))
+    z_dist = cuqi.distribution.Gaussian(
+        mean=np.zeros(3),
+        cov=np.eye(3))
+
+    # Create data distribution
+    data_dist = cuqi.distribution.Gaussian(
+        mean=model(x_dist, y_dist, z_dist), cov = 1.0)
+
+    # Create likelihood
+    likelihood = data_dist(data_dist = np.array([2,2,3]))
+
+    x_val = np.array([1, 2, 3])
+    y_val = np.array([4, 5])
+    z_val = np.array([6, 7, 8])
+
+    posterior = cuqi.distribution.JointDistribution(
+        likelihood,
+        x_dist,
+        y_dist,
+        z_dist
+    )
+
+    # Ensure correct parameter names are returned for joint distribution with likelihood
+    assert posterior.get_parameter_names() == ['x_dist', 'y_dist', 'z_dist']
+
+    assert posterior(x_dist=x_val).get_parameter_names() == ['y_dist', 'z_dist']
+    assert posterior(y_dist=y_val).get_parameter_names() == ['x_dist', 'z_dist']
+    assert posterior(z_dist=z_val).get_parameter_names() == ['x_dist', 'y_dist']
+
+    assert posterior(y_dist=y_val, z_dist=z_val).get_parameter_names() == ['x_dist']
+    assert posterior(x_dist=x_val, z_dist=z_val).get_parameter_names() == ['y_dist']
+    assert posterior(x_dist=x_val, y_dist=y_val).get_parameter_names() == ['z_dist']
+
+    assert posterior(x_dist=x_val, y_dist=y_val, z_dist=z_val).get_parameter_names() == []
+
+    joint_dist =  cuqi.distribution.JointDistribution(
+        data_dist,
+        x_dist,
+        y_dist,
+        z_dist
+    )
+
+    # Ensure correct parameter names are returned for joint distribution with data distribution
+    assert joint_dist.get_parameter_names() == ['data_dist', 'x_dist', 'y_dist', 'z_dist']
+    assert joint_dist(x_dist=x_val).get_parameter_names() == ['data_dist', 'y_dist', 'z_dist']
+    assert joint_dist(y_dist=y_val).get_parameter_names() == ['data_dist', 'x_dist', 'z_dist']
+    assert joint_dist(z_dist=z_val).get_parameter_names() == ['data_dist', 'x_dist', 'y_dist']
+    assert joint_dist(data_dist=np.array([2,2,3])).get_parameter_names() == ['x_dist', 'y_dist', 'z_dist']
+
+    assert joint_dist(x_dist=x_val, data_dist=np.array([2,2,3])).get_parameter_names() == ['y_dist', 'z_dist']
+    assert joint_dist(y_dist=y_val, data_dist=np.array([2,2,3])).get_parameter_names() == ['x_dist', 'z_dist']
+    assert joint_dist(z_dist=z_val, data_dist=np.array([2,2,3])).get_parameter_names() == ['x_dist', 'y_dist']
+    assert joint_dist(x_dist=x_val, y_dist=y_val).get_parameter_names() == ['data_dist', 'z_dist']
+    assert joint_dist(x_dist=x_val, z_dist=z_val).get_parameter_names() == ['data_dist', 'y_dist']
+    assert joint_dist(y_dist=y_val, z_dist=z_val).get_parameter_names() == ['data_dist', 'x_dist']
+
+    assert joint_dist(x_dist=x_val, y_dist=y_val, z_dist=z_val).get_parameter_names() == ['data_dist']
+    assert joint_dist(x_dist=x_val, y_dist=y_val, data_dist=np.array([2,2,3])).get_parameter_names() == ['z_dist']
+    assert joint_dist(x_dist=x_val, z_dist=z_val, data_dist=np.array([2,2,3])).get_parameter_names() == ['y_dist']
+    assert joint_dist(y_dist=y_val, z_dist=z_val, data_dist=np.array([2,2,3])).get_parameter_names() == ['x_dist']
+
+    # Ensure correct parameter names are returned for underlying likelihood
+    assert joint_dist(data_dist=np.array([2,2,3]))._likelihoods[0].get_parameter_names() == ['x_dist', 'y_dist', 'z_dist']
+
+    assert joint_dist(x_dist=x_val, data_dist=np.array([2,2,3]))._likelihoods[0].get_parameter_names() == ['y_dist', 'z_dist']
+    assert joint_dist(y_dist=y_val, data_dist=np.array([2,2,3]))._likelihoods[0].get_parameter_names() == ['x_dist', 'z_dist']
+    assert joint_dist(z_dist=z_val, data_dist=np.array([2,2,3]))._likelihoods[0].get_parameter_names() == ['x_dist', 'y_dist']
+
+    assert joint_dist(x_dist=x_val, y_dist=y_val, data_dist=np.array([2,2,3])).likelihood.get_parameter_names() == ['z_dist']
+    assert joint_dist(x_dist=x_val, z_dist=z_val, data_dist=np.array([2,2,3])).likelihood.get_parameter_names() == ['y_dist']
+    assert joint_dist(y_dist=y_val, z_dist=z_val, data_dist=np.array([2,2,3])).likelihood.get_parameter_names() == ['x_dist']