Skip to content

Commit

Permalink
added bounded generator
Browse files Browse the repository at this point in the history
  • Loading branch information
christinahedges committed Mar 18, 2024
1 parent 5a310e9 commit 69bae9b
Show file tree
Hide file tree
Showing 8 changed files with 232 additions and 39 deletions.
1 change: 1 addition & 0 deletions src/lamatrix/__init__.py
Expand Up @@ -36,6 +36,7 @@ def _META_DATA():

import numpy as np

from .bound import *
from .combine import * # noqa: E402, F401
from .models.astrophysical import * # noqa: E402, F401
from .models.gaussian import * # noqa: E402, F401
Expand Down
163 changes: 163 additions & 0 deletions src/lamatrix/bound.py
@@ -0,0 +1,163 @@
import numpy as np

from .combine import StackedIndependentGenerator

__all__ = ["BoundedGenerator"]


class BoundedGenerator(StackedIndependentGenerator):
def __init__(self, generator, bounds: list):
self.generator = generator
self.x_name = self.generator.x_name
if isinstance(bounds, slice):
self._latex_bounds = self._slice_bounds_to_latex(
bounds.start, bounds.stop, bounds.step
)
self.nbounds = (bounds.stop - bounds.start) // bounds.step

else:
self._latex_bounds = self._bounds_to_latex(bounds)
self.nbounds = len(bounds)

self.bounds = bounds
self.fit_mu = None
self.fit_sigma = None

def _bounds_to_latex(self, bounds):
bound_latex = [
f"""\\[b_{{{idx}}}(\\mathbf{{{self.x_name}}}) =
\\begin{{cases}}
\\mathbf{{{self.x_name}}}, & \\text{{if }} {bound[0]} < \\mathbf{{{self.x_name}}} \\leq {bound[1]} \\\\
0, & \\text{{otherwise}}
\\end{{cases}}\\]"""
for idx, bound in enumerate(bounds)
]
return "\n".join(bound_latex)

def _slice_bounds_to_latex(self, start, stop, step):
step_str = f"{step}{{\cdot}}" if step != 1 else ""
nbounds = (stop - start) // step
if nbounds < 4:
def_str1 = "0"
else:
def_str1 = ", ".join([f"{s}" for s in np.arange(0, 3)])
def_str = f"For $i = {def_str1}, \\ldots, {nbounds} $ define:"
bound_latex = [
f"""b_{{i}}(\\mathbf{{{self.x_name}}}) =
\\begin{{cases}}
\\mathbf{{{self.x_name}}}, & \\text{{if }} ({start} + {step_str}i) < \\mathbf{{{self.x_name}}} \\leq ({start} + {step_str}(i + 1)) \\\\
0, & \\text{{otherwise}}
\\end{{cases}}"""
]
return def_str + "\[" + "\n".join(bound_latex) + "\]"

def __repr__(self):
str1 = f"{type(self).__name__}({type(self.generator).__name__}({', '.join(list(self.arg_names))}))[n, {self.width}]"

return str1

@property
def mu(self):
return self.prior_mu if self.fit_mu is None else self.fit_mu

@property
def sigma(self):
return self.prior_sigma if self.fit_sigma is None else self.fit_sigma

@property
def prior_mu(self):
return np.hstack([self.generator.prior_mu] * self.nbounds)

@property
def prior_sigma(self):
return np.hstack([self.generator.prior_sigma] * self.nbounds)

@property
def arg_names(self):
return self.generator.arg_names

@property
def width(self):
return self.generator.width * self.nbounds

@property
def _equation(self):
eqn = ["b_{i}(" + eqn + ")" for eqn in self.generator._equation]
return eqn

@property
def equation(self):
func_signature = ", ".join([f"\\mathbf{{{a}}}" for a in self.arg_names])
n = self.nbounds
return (
f"\\[f({func_signature}) = "
+ " + ".join(
[
f"\\sum_{{i=0}}^{{{n}}} w_{{{{i}}, {coeff} }} {e}"
for coeff, e in enumerate(self._equation)
]
)
+ "\\]"
)

def to_latex(self):
return "\n".join([self._latex_bounds, self.equation, self._to_latex_table()])

def design_matrix(self, *args, **kwargs):
if not self.arg_names.issubset(set(kwargs.keys())):
raise ValueError(f"Expected {self.arg_names} to be passed.")
x = kwargs.get(self.x_name)
if isinstance(self.bounds, slice):
bounds_list = [
(a, a + self.bounds.step)
for a in np.arange(
self.bounds.start, self.bounds.stop, self.bounds.step
)
]
else:
bounds_list = self.bounds
return np.hstack(
[
self.generator.design_matrix(x=x) * ((x >= b[0]) & (x < b[1]))[:, None]
for b in bounds_list
]
)

def fit(self, *args, **kwargs):
self.fit_mu, self.fit_sigma = self._fit(*args, **kwargs)

@property
def table_properties(self):
return [
(
"w_{{{i}, {idx}}}",
(self.mu[idx], self.sigma[idx]),
(self.prior_mu[idx], self.prior_sigma[idx]),
)
for idx in range(self.width)
]

def _to_latex_table(self):
latex_table = "\\begin{table}[h!]\n\\centering\n"
latex_table += "\\begin{tabular}{|c|c|c|}\n\\hline\n"
latex_table += "Coefficient & Best Fit & Prior \\\\\\hline\n"
idx = 0
for tm in self._get_table_matter():
latex_table += tm.format(
idx=idx % self.nbounds, i=(idx - (idx % self.nbounds)) // self.nbounds
)
idx += 1
latex_table += "\\end{tabular}\n\\end{table}"
return latex_table

def __getitem__(self, key):
g = self.generator.copy()
attrs = ["fit_mu", "fit_sigma"]
for attr in attrs:
setattr(
g, attr, getattr(self, attr).reshape((self.nbounds, len(self)))[key]
)
return g

def __len__(self):
return self.width // self.nbounds
8 changes: 6 additions & 2 deletions src/lamatrix/combine.py
Expand Up @@ -188,6 +188,7 @@ def save(self, filename: str):
data_to_store["metadata"] = _META_DATA()
json.dump(data_to_store, json_file, indent=4)


class StackedDependentGenerator(StackedIndependentGenerator):

@property
Expand Down Expand Up @@ -238,6 +239,9 @@ def __getitem__(self, key):
raise AttributeError(
"Can not extract individual generators from a dependent stacked generator."
)
@property

@property
def gradient(self):
raise AttributeError("Can not create a gradient for a dependent stacked generator.")
raise AttributeError(
"Can not create a gradient for a dependent stacked generator."
)
8 changes: 3 additions & 5 deletions src/lamatrix/generator.py
Expand Up @@ -101,10 +101,8 @@ def copy(self):
return deepcopy(self)

def __repr__(self):
fit = "<fit>" if self.fit_mu is not None else ""
return (
f"{type(self).__name__}({', '.join(list(self.arg_names))})[n, {self.width}] {fit}"
)
fit = "fit" if self.fit_mu is not None else ""
return f"{type(self).__name__}({', '.join(list(self.arg_names))})[n, {self.width}] {fit}"

# def __add__(self, other):
# if isinstance(other, Generator):
Expand All @@ -128,7 +126,7 @@ def format_significant_figures(mean, error):
if error == 0:
sig_figures = 0
else:
sig_figures = -int(math.floor(math.log10(abs(error))))
sig_figures = np.max([0, -int(math.floor(math.log10(abs(error))))])

# Format mean and error to have the same number of decimal places
formatted_mean = f"{mean:.{sig_figures}f}"
Expand Down
6 changes: 3 additions & 3 deletions src/lamatrix/models/gaussian.py
Expand Up @@ -169,7 +169,7 @@ def table_properties(self):
@property
def _equation(self):
return [
f"",
f"\\mathbf{{{self.x_name}}}^0",
f"\mathbf{{{self.x_name}}}^{2}",
f"\mathbf{{{self.y_name}}}^{2}",
f"\mathbf{{{self.x_name}}}\mathbf{{{self.y_name}}}",
Expand All @@ -189,7 +189,7 @@ def to_latex(self):
)

@property
def derivative(self):
def gradient(self):
return dlnGaussian2DGenerator(
stddev_x=self.stddev_x[0],
stddev_y=self.stddev_y[0],
Expand Down Expand Up @@ -306,4 +306,4 @@ def table_properties(self):
def _equation(self):
dfdx = f"\\left(-\\frac{{1}}{{1-\\rho^2}}\\left(\\frac{{\\mathbf{{{self.x_name}}}}}{{\\sigma_x^2}} - \\rho\\frac{{\\mathbf{{{self.y_name}}}}}{{\\sigma_x\\sigma_y}}\\right)\\right)"
dfdy = f"\\left(-\\frac{{1}}{{1-\\rho^2}}\\left(\\frac{{\\mathbf{{{self.y_name}}}}}{{\\sigma_x^2}} - \\rho\\frac{{\\mathbf{{{self.x_name}}}}}{{\\sigma_x\\sigma_y}}\\right)\\right)"
return ["", dfdx, dfdy]
return [f"\\mathbf{{{self.x_name}}}^0", dfdx, dfdy]
36 changes: 25 additions & 11 deletions src/lamatrix/models/simple.py
Expand Up @@ -83,13 +83,18 @@ def _equation(self):
eqn = [
f"\mathbf{{{self.x_name}}}^{{{idx}}}" for idx in range(self.polyorder + 1)
]
eqn[0] = ""
# eqn[0] = ""
return eqn

@property
def gradient(self):
return dPolynomial1DGenerator(weights=self.mu, x_name=self.x_name, polyorder=self.polyorder, data_shape=self.data_shape, offset_prior=(self.mu[1], self.sigma[1])
)
return dPolynomial1DGenerator(
weights=self.mu,
x_name=self.x_name,
polyorder=self.polyorder,
data_shape=self.data_shape,
offset_prior=(self.mu[1], self.sigma[1]),
)


class dPolynomial1DGenerator(MathMixins, Generator):
Expand Down Expand Up @@ -152,7 +157,12 @@ def design_matrix(self, *args, **kwargs):
if not self.arg_names.issubset(set(kwargs.keys())):
raise ValueError(f"Expected {self.arg_names} to be passed.")
x = kwargs.get(self.x_name).ravel()
return np.vstack([(idx + 1) * w * x**idx for idx, w in zip(range(self.polyorder), self.weights[1:])]).T
return np.vstack(
[
(idx + 1) * w * x**idx
for idx, w in zip(range(self.polyorder), self.weights[1:])
]
).T

def fit(self, *args, **kwargs):
self.fit_mu, self.fit_sigma = self._fit(*args, **kwargs)
Expand All @@ -164,13 +174,13 @@ def offset(self):
@property
def _equation(self):
eqn = [
f"({idx + 1}v_{{\\textit{{poly}}, {idx + 1}}}\mathbf{{{self.x_name}}}^{{{idx}}})" for idx in range(self.polyorder)
f"({idx + 1}v_{{\\textit{{poly}}, {idx + 1}}}\mathbf{{{self.x_name}}}^{{{idx}}})"
for idx in range(self.polyorder)
]
eqn[0] = ""
# eqn[0] = ""
return eqn



class SinusoidGenerator(MathMixins, Generator):
def __init__(
self,
Expand Down Expand Up @@ -248,7 +258,7 @@ def frq(term):

return np.hstack(
[
"",
f"\\mathbf{{{self.x_name}}}^0",
*[
[
f"\sin({frq(idx)}\\mathbf{{{self.x_name}}})",
Expand Down Expand Up @@ -333,8 +343,12 @@ def design_matrix(self, *args, **kwargs):
[
x**0,
*[
[w1*np.cos(x * (idx + 1)), -w2*np.sin(x * (idx + 1))]
for idx, w1, w2 in zip(range(self.nterms), self.weights[1:][::2], self.weights[1:][1::2])
[w1 * np.cos(x * (idx + 1)), -w2 * np.sin(x * (idx + 1))]
for idx, w1, w2 in zip(
range(self.nterms),
self.weights[1:][::2],
self.weights[1:][1::2],
)
],
]
).T
Expand All @@ -349,7 +363,7 @@ def frq(term):

return np.hstack(
[
"",
f"\\mathbf{{{self.x_name}}}^0",
*[
[
f"v_{{\\textit{{sin}}, {idx * 2}}}\cos({frq(idx)}\\mathbf{{{self.x_name}}})",
Expand Down
4 changes: 2 additions & 2 deletions src/lamatrix/models/spline.py
Expand Up @@ -150,7 +150,7 @@ def offset(self):
@property
def _equation(self):
return [
"",
f"\\mathbf{{{self.x_name}}}^0",
*[
f"N_{{{idx},k}}(\\mathbf{{{self.x_name}}})"
for idx in np.arange(1, self.width)
Expand Down Expand Up @@ -277,6 +277,6 @@ def table_properties(self):
@property
def _equation(self):
return [
"",
f"\\mathbf{{{self.x_name}}}^0",
f"\\frac{{\\partial \\left( \\sum_{{i=1}}^{{{len(self.knots) - self.splineorder}}} v_{{\\textit{{spline}}, i}} N_{{i,{self.splineorder}}}(\\mathbf{{{self.x_name}}})\\right)}}{{\\partial \mathbf{{{self.x_name}}}}}",
]

0 comments on commit 69bae9b

Please sign in to comment.