I don't feel too comfortable with the use of object oriented programming and the related advanced topics. That's why I went through this exercise of using Python's abstract class to implement an abstract Ansatz class. (An Ansatz is basically something like a mathematical assumption, i.e. we might assume that our solution looks like a polynomial or like a logistic function).
Since different Ansätze might have different signatures, I wanted to be somewhat generic. Some Ansätze might just get a few scalar parameters, others arrays, yet others a mix of both.
All derived classes need to implement the eval function, which evaluates the Ansatz for a given input.
For all Ansätze, I want to be able to optimally fit the parameters to match some given data. That's why the abstract Ansatz class implements a fit function. I am using scipy here and the curve_fit method to fit the parameters to a function. Since curve_fit expects an array of parameters, I need to transform from my possibly nested and different Ansatz parameter types to a flat array. That's what __flatten_attributes does. The reverse of this function is __set_attributes_from_flat_array.
Since the eval method should not require to retype the parameters, but the curve_fitting requires them to be passed to the function, I created a wrapper for eval called __eval_with_parameters that sets the additional parameters as the class's attributes and then evaluates the Ansatz.
I'm interested in hearing your opinion on whether or not my approach is somewhat sensible. Are there concepts that would have simplified my implementation? Would you have done some typing differently?
Here's Ansatz.py
from abc import ABC, abstractmethod
from typing import Callable, Dict, Tuple, Optional
import matplotlib.pyplot as plt
import numpy as np
import numpy.typing as npt
from scipy.optimize import curve_fit
Array = npt.NDArray[np.float64]
Numeric = npt.ArrayLike
class Ansatz(ABC):
def __init__(self, **kwargs: Numeric) -> None:
# Store names and number of elements of **kwargs in self._dimensions
self._dimensions: Dict[str, Tuple[int]] = {}
# Set instance attributes
for key, value in kwargs.items():
if np.isscalar(value):
setattr(self, key, value)
self._dimensions[key] = (1,)
else:
value_as_array: npt.NDArray[np.float64] = np.array(value)
setattr(self, key, value_as_array)
self._dimensions[key] = value_as_array.shape
# Vectorized version of the Ansatz
self.__eval_vectorized: Callable[[Array], Array] = np.vectorize(self.eval)
@abstractmethod
def eval(self, x: float) -> float:
"""Evaluate the Ansatz.
Parameters
----------
x : float
Where to evaluate the Ansatz.
Returns
-------
float
The evaluated Ansatz.
"""
pass
def fit(self, x: Array, y: Array, plot: bool = False, *args, **kwargs) -> None:
"""Fit the Ansatz to the data. Overwrite the instance attributes.
Sets the instance attributes to the fitted parameters unless there is
a RuntimeErrror. In that case, revert the instance attributes to the
original values.
Parameters
----------
x : Array
x-values of the data used for fitting.
y : Array
y-values of the data used for fitting.
plot : bool, optional
Whether or not to plot the data and fit, by default False
"""
p0 = self.__flatten_attributes()
try:
p_opt, _ = curve_fit(
self.__eval_with_parameters, x, y, p0=p0, *args, **kwargs
)
self.__set_attributes_from_flat_array(p_opt)
if plot:
t = np.linspace(np.min(x), np.max(x), 100)
_, ax = plt.subplots()
ax.plot(x, y, "o", label="data")
ax.plot(t, self.__eval_vectorized(t), "-", label="fit")
ax.set_title(f"{self.__class__.__name__} Ansatz")
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.legend()
plt.show()
except RuntimeError as e:
print(e)
self.__set_attributes_from_flat_array(p0)
def __eval_with_parameters(self, x: Array, *args: float) -> Array:
"""Evaluate the Ansatz with the given parameters.
Changes the instance attributes.
Parameters
----------
x : Array
Where to evaluate the Ansatz.
*args : float
The parameters used to temporarily set the instance attributes to.
Returns
-------
Array
The evaluated Ansatz.
"""
self.__set_attributes_from_flat_array(np.array(args))
return self.__eval_vectorized(x)
def __flatten_attributes(self) -> Array:
"""Return all instance attributes as a flattened array.
Returns
-------
Array
Flat array with all instance attributes.
"""
flat_attributes = np.empty(0)
for key, _ in self._dimensions.items():
attr = getattr(self, key)
if not np.isscalar(attr):
attr = attr.flatten()
flat_attributes = np.append(flat_attributes, attr)
return flat_attributes
def __set_attributes_from_flat_array(self, flat_attributes: Array) -> None:
"""Sets the instance attributes from a flattened array.
Parameters
----------
flat_attributes : Numeric
Flat array with attributes to set.
"""
counter = 0
for key, shape in self._dimensions.items():
n_elements = int(np.prod(shape))
attr = flat_attributes[counter : counter + n_elements].reshape(shape)
setattr(self, key, attr)
counter += n_elements
And here are some derived implementations with a few runnable examples (I omitted the comments that describe the different Ansätze).
from typing import Optional
import numpy as np
import numpy.typing as npt
from .Ansatz import Ansatz
Array = npt.NDArray[np.float64]
Numeric = npt.ArrayLike
class Poly(Ansatz):
coefficients: Array
def __init__(
self, order: Optional[int], coefficients: Optional[Numeric] = None
) -> None:
if coefficients is None:
if order is not None:
coefficients = np.zeros(order + 1)
else:
raise ValueError("Either order or coefficients must be given.")
if coefficients is not None and order is not None:
if np.size(coefficients) != order + 1:
raise ValueError(
f"Order of polynomial ({order}) +1 does not match number of coefficients ({np.size(coefficients)})"
)
super().__init__(coefficients=coefficients)
def eval(self, x: float) -> float:
powers_of_x = [x**i for i in range(len(self.coefficients))]
return np.dot(self.coefficients, powers_of_x)
class Exponential(Ansatz):
y0: float
rate: float
def __init__(self, y0: float = 1.0, rate: float = 1.0):
super().__init__(y0=y0, rate=rate)
def eval(self, x: float) -> float:
return self.y0 * np.exp(self.rate * x)
class Rational(Ansatz):
offset: float
enumerator: Array
denominator: Array
def __init__(
self, offset: float = 0.0, enumerator: Numeric = 1.0, denominator: Numeric = 1.0
):
assert np.size(enumerator) == np.size(denominator)
super().__init__(offset=offset, enumerator=enumerator, denominator=denominator)
def eval(self, x: float) -> float:
powers_of_x = [x**i for i in range(len(self.enumerator))]
return (
np.dot(self.enumerator, powers_of_x) / np.dot(self.denominator, powers_of_x)
) + self.offset
if __name__ == "__main__":
P = Poly(order=2, coefficients=[1, 2, 3])
E = Exponential(y0=1.0, rate=2.0)
R = Rational(offset=1.0, enumerator=[1, 2], denominator=[1, 1])
x = np.linspace(0, 10, 100)
y = np.random.randn(x.shape[0]) + x**2 - x**3 / 10
P.fit(x, y, plot=True)
E.fit(x, y, plot=True)
R.fit(x, y, plot=True)
Thanks in advance for all feedback!
