From 01065de82b170108ce7a8d43deee6082c7e692a5 Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Mon, 2 Dec 2024 09:39:27 -0700 Subject: [PATCH] coverage --- named_arrays/_functions/functions.py | 27 +++++++++++++-------------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/named_arrays/_functions/functions.py b/named_arrays/_functions/functions.py index 36b3e7b..d517a48 100644 --- a/named_arrays/_functions/functions.py +++ b/named_arrays/_functions/functions.py @@ -842,28 +842,27 @@ def design_matrix( design_matrix = {} if isinstance(inputs, na.AbstractScalar): - inputs = na.CartesianNdVectorArray({'dummy': inputs}) - - cartesian_nd = inputs.cartesian_nd.broadcasted - # grab subset of components involved in polynomial fit - if self.components_polynomial: - cartesian_nd = na.CartesianNdVectorArray( - { - k: cartesian_nd.components[k] - for k in cartesian_nd.components.keys() & self.components_polynomial - } - ) + inputs = na.CartesianNdVectorArray({"dummy": inputs}) + inputs = inputs.cartesian_nd.broadcasted.components + + components = self.components_polynomial + + if components is None: + components = tuple(inputs) + else: + components = (components,) if isinstance(components, str) else components + + inputs = {c: inputs[c] for c in components} for i in range(self.degree + 1): combinations = itertools.combinations_with_replacement( - cartesian_nd.components, i + inputs, i ) for combination in combinations: key = "*".join(combination) design_matrix[key] = 1 for k in combination: - design_matrix[key] = design_matrix[key] * cartesian_nd.components[k] - + design_matrix[key] = design_matrix[key] * inputs[k] design_matrix = na.CartesianNdVectorArray(design_matrix)