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

Add UP rule ruff #635

Merged
merged 3 commits into from
Apr 16, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
4 changes: 2 additions & 2 deletions docs/source/notebooks/clv/sBG.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -1169,8 +1169,8 @@
"source": [
"weeks = xr.DataArray(np.arange(1, 14), dims=[\"weeks\"])\n",
"\n",
"survival_perc_highend = 100-100*((churn_highend < weeks))\n",
"survival_perc_regular = 100-100*((churn_regular < weeks))"
"survival_perc_highend = 100-100*(churn_highend < weeks)\n",
"survival_perc_regular = 100-100*(churn_regular < weeks)"
]
},
{
Expand Down
13 changes: 7 additions & 6 deletions pymc_marketing/clv/models/basic.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
import json
import warnings
from collections.abc import Sequence
from pathlib import Path
from typing import Dict, Optional, Sequence, cast
from typing import cast

import arviz as az
import pandas as pd
Expand All @@ -21,8 +22,8 @@ def __init__(
self,
data: pd.DataFrame,
*,
model_config: Optional[Dict] = None,
sampler_config: Optional[Dict] = None,
model_config: dict | None = None,
sampler_config: dict | None = None,
):
super().__init__(model_config, sampler_config)
self.data = data
Expand Down Expand Up @@ -211,18 +212,18 @@ def thin_fit_result(self, keep_every: int):
return type(self)._build_with_idata(new_idata)

@staticmethod
def _create_distribution(dist: Dict, shape=()):
def _create_distribution(dist: dict, shape=()):
try:
return getattr(pm, dist["dist"]).dist(**dist.get("kwargs", {}), shape=shape)
except AttributeError:
raise ValueError(f"Distribution {dist['dist']} does not exist in PyMC")

@property
def default_sampler_config(self) -> Dict:
def default_sampler_config(self) -> dict:
return {}

@property
def _serializable_model_config(self) -> Dict:
def _serializable_model_config(self) -> dict:
return self.model_config

@property
Expand Down
34 changes: 17 additions & 17 deletions pymc_marketing/clv/models/beta_geo.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Dict, Optional, Sequence, Union
from collections.abc import Sequence

import numpy as np
import pandas as pd
Expand Down Expand Up @@ -117,8 +117,8 @@ class BetaGeoModel(CLVModel):
def __init__(
self,
data: pd.DataFrame,
model_config: Optional[Dict] = None,
sampler_config: Optional[Dict] = None,
model_config: dict | None = None,
sampler_config: dict | None = None,
):
self._validate_cols(
data,
Expand All @@ -132,7 +132,7 @@ def __init__(
)

@property
def default_model_config(self) -> Dict[str, Dict]:
def default_model_config(self) -> dict[str, dict]:
return {
"a_prior": {"dist": "HalfFlat", "kwargs": {}},
"b_prior": {"dist": "HalfFlat", "kwargs": {}},
Expand Down Expand Up @@ -211,11 +211,11 @@ def _unload_params(self):
# taken from https://lifetimes.readthedocs.io/en/latest/lifetimes.fitters.html
def expected_num_purchases(
self,
customer_id: Union[np.ndarray, pd.Series],
t: Union[np.ndarray, pd.Series, TensorVariable],
frequency: Union[np.ndarray, pd.Series, TensorVariable],
recency: Union[np.ndarray, pd.Series, TensorVariable],
T: Union[np.ndarray, pd.Series, TensorVariable],
customer_id: np.ndarray | pd.Series,
t: np.ndarray | pd.Series | TensorVariable,
frequency: np.ndarray | pd.Series | TensorVariable,
recency: np.ndarray | pd.Series | TensorVariable,
T: np.ndarray | pd.Series | TensorVariable,
) -> xr.DataArray:
r"""
Given a purchase history/profile of :math:`x` and :math:`t_x` for an individual
Expand Down Expand Up @@ -273,10 +273,10 @@ def expected_num_purchases(

def expected_probability_alive(
self,
customer_id: Union[np.ndarray, pd.Series],
frequency: Union[np.ndarray, pd.Series],
recency: Union[np.ndarray, pd.Series],
T: Union[np.ndarray, pd.Series],
customer_id: np.ndarray | pd.Series,
frequency: np.ndarray | pd.Series,
recency: np.ndarray | pd.Series,
T: np.ndarray | pd.Series,
) -> xr.DataArray:
r"""
Posterior expected value of the probability of being alive at time T. The
Expand Down Expand Up @@ -310,7 +310,7 @@ def expected_probability_alive(

def expected_num_purchases_new_customer(
self,
t: Union[np.ndarray, pd.Series],
t: np.ndarray | pd.Series,
):
r"""
Posterior expected number of purchases for any interval of length :math:`t`. See
Expand Down Expand Up @@ -345,7 +345,7 @@ def expected_num_purchases_new_customer(

def _distribution_new_customers(
self,
random_seed: Optional[RandomState] = None,
random_seed: RandomState | None = None,
var_names: Sequence[str] = ("population_dropout", "population_purchase_rate"),
) -> xr.Dataset:
with pm.Model():
Expand All @@ -372,7 +372,7 @@ def _distribution_new_customers(

def distribution_new_customer_dropout(
self,
random_seed: Optional[RandomState] = None,
random_seed: RandomState | None = None,
) -> xr.Dataset:
"""Sample the Beta distribution for the population-level dropout rate.

Expand All @@ -396,7 +396,7 @@ def distribution_new_customer_dropout(

def distribution_new_customer_purchase_rate(
self,
random_seed: Optional[RandomState] = None,
random_seed: RandomState | None = None,
) -> xr.Dataset:
"""Sample the Gamma distribution for the population-level purchase rate.

Expand Down
58 changes: 28 additions & 30 deletions pymc_marketing/clv/models/gamma_gamma.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
from typing import Dict, Optional, Union

import numpy as np
import pandas as pd
import pymc as pm
Expand All @@ -15,10 +13,10 @@
class BaseGammaGammaModel(CLVModel):
def distribution_customer_spend(
self,
customer_id: Union[np.ndarray, pd.Series],
mean_transaction_value: Union[np.ndarray, pd.Series, TensorVariable],
frequency: Union[np.ndarray, pd.Series, TensorVariable],
random_seed: Optional[RandomState] = None,
customer_id: np.ndarray | pd.Series,
mean_transaction_value: np.ndarray | pd.Series | TensorVariable,
frequency: np.ndarray | pd.Series | TensorVariable,
random_seed: RandomState | None = None,
) -> xarray.DataArray:
"""Posterior distribution of transaction value per customer"""

Expand All @@ -44,9 +42,9 @@ def distribution_customer_spend(

def expected_customer_spend(
self,
customer_id: Union[np.ndarray, pd.Series],
mean_transaction_value: Union[np.ndarray, pd.Series],
frequency: Union[np.ndarray, pd.Series],
customer_id: np.ndarray | pd.Series,
mean_transaction_value: np.ndarray | pd.Series,
frequency: np.ndarray | pd.Series,
) -> xarray.DataArray:
"""Expected transaction value per customer

Expand Down Expand Up @@ -107,11 +105,11 @@ def expected_new_customer_spend(self) -> xarray.DataArray:
def expected_customer_lifetime_value(
self,
transaction_model: CLVModel,
customer_id: Union[np.ndarray, pd.Series],
mean_transaction_value: Union[np.ndarray, pd.Series],
frequency: Union[np.ndarray, pd.Series],
recency: Union[np.ndarray, pd.Series],
T: Union[np.ndarray, pd.Series],
customer_id: np.ndarray | pd.Series,
mean_transaction_value: np.ndarray | pd.Series,
frequency: np.ndarray | pd.Series,
recency: np.ndarray | pd.Series,
T: np.ndarray | pd.Series,
time: int = 12,
discount_rate: float = 0.01,
freq: str = "D",
Expand Down Expand Up @@ -226,8 +224,8 @@ class GammaGammaModel(BaseGammaGammaModel):
def __init__(
self,
data: pd.DataFrame,
model_config: Optional[Dict] = None,
sampler_config: Optional[Dict] = None,
model_config: dict | None = None,
sampler_config: dict | None = None,
):
self._validate_cols(
data,
Expand All @@ -239,7 +237,7 @@ def __init__(
)

@property
def default_model_config(self) -> Dict:
def default_model_config(self) -> dict:
return {
"p_prior": {"dist": "HalfFlat", "kwargs": {}},
"q_prior": {"dist": "HalfFlat", "kwargs": {}},
Expand Down Expand Up @@ -362,8 +360,8 @@ class GammaGammaModelIndividual(BaseGammaGammaModel):
def __init__(
self,
data: pd.DataFrame,
model_config: Optional[Dict] = None,
sampler_config: Optional[Dict] = None,
model_config: dict | None = None,
sampler_config: dict | None = None,
):
self._validate_cols(
data, required_cols=["customer_id", "individual_transaction_value"]
Expand All @@ -373,7 +371,7 @@ def __init__(
)

@property
def default_model_config(self) -> Dict:
def default_model_config(self) -> dict:
return {
"p_prior": {"dist": "HalfFlat", "kwargs": {}},
"q_prior": {"dist": "HalfFlat", "kwargs": {}},
Expand Down Expand Up @@ -419,9 +417,9 @@ def _summarize_mean_data(self, customer_id, individual_transaction_value):

def distribution_customer_spend( # type: ignore [override]
self,
customer_id: Union[np.ndarray, pd.Series],
individual_transaction_value: Union[np.ndarray, pd.Series, TensorVariable],
random_seed: Optional[RandomState] = None,
customer_id: np.ndarray | pd.Series,
individual_transaction_value: np.ndarray | pd.Series | TensorVariable,
random_seed: RandomState | None = None,
) -> xarray.DataArray:
"""Return distribution of transaction value per customer"""

Expand All @@ -438,9 +436,9 @@ def distribution_customer_spend( # type: ignore [override]

def expected_customer_spend(
self,
customer_id: Union[np.ndarray, pd.Series],
individual_transaction_value: Union[np.ndarray, pd.Series, TensorVariable],
random_seed: Optional[RandomState] = None,
customer_id: np.ndarray | pd.Series,
individual_transaction_value: np.ndarray | pd.Series | TensorVariable,
random_seed: RandomState | None = None,
) -> xarray.DataArray:
"""Return expected transaction value per customer"""

Expand All @@ -458,10 +456,10 @@ def expected_customer_spend(
def expected_customer_lifetime_value( # type: ignore [override]
self,
transaction_model: CLVModel,
customer_id: Union[np.ndarray, pd.Series],
individual_transaction_value: Union[np.ndarray, pd.Series, TensorVariable],
recency: Union[np.ndarray, pd.Series],
T: Union[np.ndarray, pd.Series],
customer_id: np.ndarray | pd.Series,
individual_transaction_value: np.ndarray | pd.Series | TensorVariable,
recency: np.ndarray | pd.Series,
T: np.ndarray | pd.Series,
time: int = 12,
discount_rate: float = 0.01,
freq: str = "D",
Expand Down
Loading
Loading