Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Internal refactor #804

Merged
merged 21 commits into from
Jul 4, 2024
Merged
Show file tree
Hide file tree
Changes from 18 commits
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions .pylintrc
Original file line number Diff line number Diff line change
Expand Up @@ -384,8 +384,7 @@ missing-member-max-choices=1

# List of note tags to take in consideration, separated by a comma.
notes=FIXME,
XXX,
TODO
XXX


[SPELLING]
Expand Down
48 changes: 22 additions & 26 deletions bambi/backend/model_components.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
from bambi.backend.utils import get_distribution_from_prior
from bambi.families.multivariate import MultivariateFamily
from bambi.families.univariate import Categorical, Cumulative, StoppingRatio
from bambi.utils import get_aliased_name


ORDINAL_FAMILIES = (Cumulative, StoppingRatio)
Expand All @@ -18,18 +17,14 @@ def __init__(self, component):
self.output = 0

def build(self, pymc_backend, bmb_model):
response_aliased_name = get_aliased_name(bmb_model.response_component.response_term)
if self.component.alias:
label = self.component.alias
else:
label = f"{response_aliased_name}_{self.component.name}"
label = self.component.alias if self.component.alias else self.component.name

# NOTE: This could be handled in a different manner in the future, only applies to
# thresholds and assumes we always do it when using ordinal families.
extra_args = {}
if isinstance(bmb_model.family, ORDINAL_FAMILIES):
threshold_dim = label + "_dim"
threshold_values = np.arange(len(bmb_model.response_component.response_term.levels) - 1)
threshold_values = np.arange(len(bmb_model.response_component.term.levels) - 1)
extra_args["dims"] = threshold_dim
pymc_backend.model.add_coords({threshold_dim: threshold_values})

Expand Down Expand Up @@ -58,7 +53,7 @@ def build(self, pymc_backend, bmb_model):
self.build_intercept(bmb_model)
self.build_offsets()
self.build_common_terms(pymc_backend, bmb_model)
self.build_hsgp_terms(pymc_backend, bmb_model)
self.build_hsgp_terms(pymc_backend)
self.build_group_specific_terms(pymc_backend, bmb_model)

def build_intercept(self, bmb_model):
Expand All @@ -78,7 +73,7 @@ def build_common_terms(self, pymc_backend, bmb_model):
"""Add common (fixed) terms to the PyMC model.

We have linear predictors of the form 'X @ b + Z @ u'.
This creates the 'b' parameter vector in PyMC, computes `X @ b`, and adds it to ``self.mu``.
This creates the 'b' parameter vector in PyMC, computes `X @ b`, and adds it to `self.mu`.

Parameters
----------
Expand Down Expand Up @@ -114,25 +109,25 @@ def build_common_terms(self, pymc_backend, bmb_model):
# Add term to linear predictor
self.output += pt.dot(data, coefs)

def build_hsgp_terms(self, pymc_backend, bmb_model):
def build_hsgp_terms(self, pymc_backend):
"""Add HSGP (Hilbert-Space Gaussian Process approximation) terms to the PyMC model.

The linear predictor 'X @ b + Z @ u' can be augmented with non-parametric HSGP terms
'f(x)'. This creates the 'f(x)' and adds it ``self.output``.
'f(x)'. This creates the 'f(x)' and adds it `self.output`.
"""
for term in self.component.hsgp_terms.values():
hsgp_term = HSGPTerm(term)
for name, values in hsgp_term.coords.items():
if name not in pymc_backend.model.coords:
pymc_backend.model.add_coords({name: values})
self.output += hsgp_term.build(bmb_model)
self.output += hsgp_term.build()

def build_group_specific_terms(self, pymc_backend, bmb_model):
"""Add group-specific (random or varying) terms to the PyMC model.

We have linear predictors of the form 'X @ b + Z @ u'.
This creates the 'u' parameter vector in PyMC, computes `Z @ u`, and adds it to
``self.output``.
`self.output`.
"""
for term in self.component.group_specific_terms.values():
group_specific_term = GroupSpecificTerm(term, bmb_model.noncentered)
Expand All @@ -156,22 +151,23 @@ def build_group_specific_terms(self, pymc_backend, bmb_model):
else:
self.output += coef * predictor

def build_response(self, pymc_backend, bmb_model):
# Extract the response term from the Bambi family
response_term = bmb_model.response_component.response_term

# Create and build the response term
response_term = ResponseTerm(response_term, bmb_model.family)
response_term.build(pymc_backend, bmb_model)

def add_response_coords(self, pymc_backend, bmb_model):
response_term = bmb_model.response_component.response_term
response_name = get_aliased_name(response_term)
dim_name = f"{response_name}_obs"
response_term = bmb_model.response_component.term
dim_name = "__obs__"
dim_value = np.arange(response_term.shape[0])
pymc_backend.model.add_coords({dim_name: dim_value})


class ResponseComponent:
def __init__(self, component):
self.component = component

def build(self, pymc_backend, bmb_model):
# Create and build the response term
response_term = ResponseTerm(self.component.term, bmb_model.family)
response_term.build(pymc_backend, bmb_model)


# # NOTE: Here for historical reasons, not supposed to work now at least for now
# def add_lkj(backend, terms, eta=1):
# """Add correlated prior for group-specific effects.
Expand All @@ -183,15 +179,15 @@ def add_response_coords(self, pymc_backend, bmb_model):
# Parameters
# ----------
# terms: list
# A list of terms that share a common grouper (i.e. ``1|Group`` and ``Variable|Group`` in
# A list of terms that share a common grouper (i.e. `1|Group` and `Variable|Group` in
# formula notation).
# eta: num
# The value for the eta parameter in the LKJ distribution.

# Parameters
# ----------
# mu
# The contribution to the linear predictor of the roup-specific terms in ``terms``.
# The contribution to the linear predictor of the roup-specific terms in `terms`.
# """

# # Parameters
Expand Down
Loading
Loading