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.

Sign up for GitHub

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

Refactor conversion to InferenceData #44

Merged
merged 12 commits into from
May 28, 2021
Merged

Refactor conversion to InferenceData #44

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
c1d7af3
Remove dask.delayed decorator from single feat fit
gibsramen May 17, 2021
06470ff
Preliminary restructure of inference conversion
gibsramen May 17, 2021
e47b029
Add auto-conversion to inference of // models
gibsramen May 17, 2021
f5cf57d
Remove unwanted vars when auto-converting to inf
gibsramen May 17, 2021
648573d
Update default model specifications
gibsramen May 17, 2021
294f726
Restructure again with comments
gibsramen May 19, 2021
0eebfbf
Remove dask cluster & jobs args
gibsramen May 26, 2021
f5f7e01
Concatenate inferences in to_inference_object
gibsramen May 26, 2021
039e0f9
Raise error if attempted fit before compilation
gibsramen May 27, 2021
b20f1af
Setup changes
gibsramen May 27, 2021
ddc1ca9
Update custom model docs page
gibsramen May 27, 2021
53bd058
Update parallelization docs
gibsramen May 27, 2021
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
150 changes: 51 additions & 99 deletions birdman/default_models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,7 @@
import os
from pkg_resources import resource_filename

import arviz as az
import biom
import dask_jobqueue
import numpy as np
import pandas as pd

Expand Down Expand Up @@ -108,51 +106,26 @@ def __init__(
}
self.add_parameters(param_dict)

def to_inference_object(
self,
dask_cluster: dask_jobqueue.JobQueueCluster = None,
jobs: int = 4
) -> az.InferenceData:
"""Convert fitted Stan model into ``arviz`` InferenceData object.

:param dask_cluster: Dask jobqueue to run parallel jobs (optional)
:type dask_cluster: dask_jobqueue.JobQueueCluster, optional

:param jobs: Number of jobs to run in parallel, defaults to 4
:type jobs: int

:returns: ``arviz`` InferenceData object with selected values
:rtype: az.InferenceData
"""
dims = {
"beta": ["covariate", "feature"],
"phi": ["feature"],
"log_lhood": ["tbl_sample", "feature"],
"y_predict": ["tbl_sample", "feature"]
}
coords = {
"covariate": self.colnames,
"feature": self.feature_names,
"tbl_sample": self.sample_names
}

# TODO: May want to allow not passing PP/LL/OD in the future
args = dict()
if self.parallelize_across == "chains":
args["alr_params"] = ["beta"]
else:
args["dask_cluster"] = dask_cluster
args["jobs"] = jobs
inf = super().to_inference_object(
self.specify_model(
params=["beta", "phi"],
dims=dims,
coords=coords,
posterior_predictive="y_predict",
log_likelihood="log_lhood",
dims={
"beta": ["covariate", "feature"],
"phi": ["feature"],
"log_lhood": ["tbl_sample", "feature"],
"y_predict": ["tbl_sample", "feature"]
},
coords={
"covariate": self.colnames,
"feature": self.feature_names,
"tbl_sample": self.sample_names
},
include_observed_data=True,
**args
posterior_predictive="y_predict",
log_likelihood="log_lhood"
)
return inf

if self.parallelize_across == "chains":
self.specifications["alr_params"] = ["beta"]


class NegativeBinomialLME(RegressionModel):
Expand Down Expand Up @@ -259,37 +232,28 @@ def __init__(
}
self.add_parameters(param_dict)

def to_inference_object(self) -> az.InferenceData:
"""Convert fitted Stan model into ``arviz`` InferenceData object.

:returns: ``arviz`` InferenceData object with selected values
:rtype: az.InferenceData
"""
dims = {
"beta": ["covariate", "feature"],
"phi": ["feature"],
"subj_int": ["group", "feature"],
"log_lhood": ["tbl_sample", "feature"],
"y_predict": ["tbl_sample", "feature"]
}
coords = {
"covariate": self.colnames,
"feature": self.feature_names,
"tbl_sample": self.sample_names,
"group": self.groups
}

# TODO: May want to allow not passing PP/LL/OD in the future
inf = super().to_inference_object(
params=["beta", "phi"],
dims=dims,
coords=coords,
posterior_predictive="y_predict",
log_likelihood="log_lhood",
alr_params=["beta", "subj_int"],
self.specify_model(
params=["beta", "phi", "subj_int"],
dims={
"beta": ["covariate", "feature"],
"phi": ["feature"],
"subj_int": ["group", "feature"],
"log_lhood": ["tbl_sample", "feature"],
"y_predict": ["tbl_sample", "feature"]
},
coords={
"covariate": self.colnames,
"feature": self.feature_names,
"tbl_sample": self.sample_names,
"group": self.groups
},
include_observed_data=True,
posterior_predictive="y_predict",
log_likelihood="log_lhood"
)
return inf

if self.parallelize_across == "chains":
self.specifications["alr_params"] = ["beta", "subj_int"]


class Multinomial(RegressionModel):
Expand Down Expand Up @@ -363,31 +327,19 @@ def __init__(
}
self.add_parameters(param_dict)

def to_inference_object(self) -> az.InferenceData:
"""Convert fitted Stan model into ``arviz`` InferenceData object.

:returns: ``arviz`` InferenceData object with selected values
:rtype: az.InferenceData
"""
dims = {
"beta": ["covariate", "feature"],
"log_lhood": ["tbl_sample"],
"y_predict": ["tbl_sample", "feature"]
}
coords = {
"covariate": self.colnames,
"feature": self.feature_names,
"tbl_sample": self.sample_names
}

# TODO: May want to allow not passing PP/LL/OD in the future
inf = super().to_inference_object(
params=["beta", "phi"],
dims=dims,
coords=coords,
alr_params=["beta"],
posterior_predictive="y_predict",
log_likelihood="log_lhood",
self.specify_model(
params=["beta"],
dims={
"beta": ["covariate", "feature"],
"log_lhood": ["tbl_sample"],
"y_predict": ["tbl_sample", "feature"]
},
coords={
"covariate": self.colnames,
"feature": self.feature_names,
"tbl_sample": self.sample_names,
},
include_observed_data=True,
posterior_predictive="y_predict",
log_likelihood="log_lhood"
)
return inf
Loading