deepecho 0.5.0.dev0__tar.gz → 0.6.0.dev0__tar.gz

deepecho-0.6.0.dev0/PKG-INFO

@@ -0,0 +1,241 @@
+Metadata-Version: 2.1
+Name: deepecho
+Version: 0.6.0.dev0
+Summary: Create sequential synthetic data of mixed types using a GAN.
+Author-email: "DataCebo, Inc." <info@sdv.dev>
+License: BSL-1.1
+Project-URL: Source Code, https://github.com/sdv-dev/Deepecho/
+Project-URL: Issue Tracker, https://github.com/sdv-dev/Deepecho/issues
+Project-URL: Twitter, https://twitter.com/sdv_dev
+Project-URL: Chat, https://bit.ly/sdv-slack-invite
+Keywords: deepecho,DeepEcho
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: Free for non-commercial use
+Classifier: Natural Language :: English
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: <3.13,>=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.20.0; python_version < "3.10"
+Requires-Dist: numpy>=1.23.3; python_version >= "3.10" and python_version < "3.12"
+Requires-Dist: numpy>=1.26.0; python_version >= "3.12"
+Requires-Dist: pandas>=1.1.3; python_version < "3.10"
+Requires-Dist: pandas>=1.3.4; python_version >= "3.10" and python_version < "3.11"
+Requires-Dist: pandas>=1.5.0; python_version >= "3.11"
+Requires-Dist: torch>=1.8.0; python_version < "3.10"
+Requires-Dist: torch>=1.11.0; python_version >= "3.10" and python_version < "3.11"
+Requires-Dist: torch>=2.0.0; python_version >= "3.11" and python_version < "3.12"
+Requires-Dist: torch>=2.2.0; python_version >= "3.12"
+Requires-Dist: tqdm>=4.29
+Provides-Extra: test
+Requires-Dist: pytest>=3.4.2; extra == "test"
+Requires-Dist: pytest-cov>=2.6.0; extra == "test"
+Requires-Dist: pytest-rerunfailures<15,>=10.3; extra == "test"
+Requires-Dist: jupyter<2,>=1.0.0; extra == "test"
+Requires-Dist: rundoc<0.5,>=0.4.3; extra == "test"
+Requires-Dist: pytest-runner>=2.11.1; extra == "test"
+Requires-Dist: tomli<3,>=2.0.0; extra == "test"
+Provides-Extra: dev
+Requires-Dist: deepecho[test]; extra == "dev"
+Requires-Dist: setuptools<70; extra == "dev"
+Requires-Dist: build<2,>=1.0.0; extra == "dev"
+Requires-Dist: bump-my-version<1,>=0.18.3; extra == "dev"
+Requires-Dist: pip>=9.0.1; extra == "dev"
+Requires-Dist: watchdog<5,>=1.0.1; extra == "dev"
+Requires-Dist: ruff<1,>=0.3.2; extra == "dev"
+Requires-Dist: twine<4,>=1.10.0; extra == "dev"
+Requires-Dist: wheel>=0.30.0; extra == "dev"
+Requires-Dist: coverage<6,>=4.5.1; extra == "dev"
+Requires-Dist: tox<4,>=2.9.1; extra == "dev"
+Requires-Dist: invoke; extra == "dev"
+
+<div align="center">
+<br/>
+<p align="center">
+    <i>This repository is part of <a href="https://sdv.dev">The Synthetic Data Vault Project</a>, a project from <a href="https://datacebo.com">DataCebo</a>.</i>
+</p>
+
+[![Development Status](https://img.shields.io/badge/Development%20Status-2%20--%20Pre--Alpha-yellow)](https://pypi.org/search/?c=Development+Status+%3A%3A+2+-+Pre-Alpha)
+[![PyPi Shield](https://img.shields.io/pypi/v/deepecho.svg)](https://pypi.python.org/pypi/deepecho)
+[![Tests](https://github.com/sdv-dev/DeepEcho/workflows/Run%20Tests/badge.svg)](https://github.com/sdv-dev/DeepEcho/actions?query=workflow%3A%22Run+Tests%22+branch%3Amain)
+[![Downloads](https://pepy.tech/badge/deepecho)](https://pepy.tech/project/deepecho)
+[![Coverage Status](https://codecov.io/gh/sdv-dev/DeepEcho/branch/main/graph/badge.svg)](https://codecov.io/gh/sdv-dev/DeepEcho)
+[![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/sdv-dev/DeepEcho/main?filepath=tutorials/timeseries_data)
+[![Slack](https://img.shields.io/badge/Slack%20Workspace-Join%20now!-36C5F0?logo=slack)](https://bit.ly/sdv-slack-invite)
+
+<div align="left">
+<br/>
+<p align="center">
+<a href="https://github.com/sdv-dev/DeepEcho">
+<img align="center" width=40% src="https://github.com/sdv-dev/SDV/blob/stable/docs/images/DeepEcho-DataCebo.png"></img>
+</a>
+</p>
+</div>
+
+</div>
+
+# Overview
+
+**DeepEcho** is a **Synthetic Data Generation** Python library for **mixed-type**, **multivariate
+time series**. It provides:
+
+1. Multiple models based both on **classical statistical modeling** of time series and the latest
+   in **Deep Learning** techniques.
+2. A robust [benchmarking framework](https://github.com/sdv-dev/SDGym) for evaluating these methods
+   on multiple datasets and with multiple metrics.
+3. Ability for **Machine Learning researchers** to submit new methods following our `model` and
+   `sample` API and get evaluated.
+
+| Important Links                               |                                                                      |
+| --------------------------------------------- | -------------------------------------------------------------------- |
+| :computer: **[Website]**                      | Check out the SDV Website for more information about the project.    |
+| :orange_book: **[SDV Blog]**                  | Regular publshing of useful content about Synthetic Data Generation. |
+| :book: **[Documentation]**                    | Quickstarts, User and Development Guides, and API Reference.         |
+| :octocat: **[Repository]**                    | The link to the Github Repository of this library.                   |
+| :keyboard: **[Development Status]**           | This software is in its Pre-Alpha stage.                             |
+| [![][Slack Logo] **Community**][Community]    | Join our Slack Workspace for announcements and discussions.          |
+| [![][MyBinder Logo] **Tutorials**][Tutorials] | Run the SDV Tutorials in a Binder environment.                       |
+
+[Website]: https://sdv.dev
+[SDV Blog]: https://sdv.dev/blog
+[Documentation]: https://sdv.dev/SDV
+[Repository]: https://github.com/sdv-dev/DeepEcho
+[License]: https://github.com/sdv-dev/DeepEcho/blob/main/LICENSE
+[Development Status]: https://pypi.org/search/?c=Development+Status+%3A%3A+2+-+Pre-Alpha
+[Slack Logo]: https://github.com/sdv-dev/SDV/blob/stable/docs/images/slack.png
+[Community]: https://bit.ly/sdv-slack-invite
+[MyBinder Logo]: https://github.com/sdv-dev/SDV/blob/stable/docs/images/mybinder.png
+[Tutorials]: https://mybinder.org/v2/gh/sdv-dev/DeepEcho/main?filepath=tutorials
+
+# Install
+
+**DeepEcho** is part of the **SDV** project and is automatically installed alongside it. For
+details about this process please visit the [SDV Installation Guide](
+https://sdv.dev/SDV/getting_started/install.html)
+
+Optionally, **DeepEcho** can also be installed as a standalone library using the following commands:
+
+**Using `pip`:**
+
+```bash
+pip install deepecho
+```
+
+**Using `conda`:**
+
+```bash
+conda install -c pytorch -c conda-forge deepecho
+```
+
+For more installation options please visit the [DeepEcho installation Guide](INSTALL.md)
+
+# Quickstart
+
+**DeepEcho** is included as part of [SDV](https://sdv.dev/SDV) to model and sample synthetic
+time series. In most cases, usage through SDV is recommeded, since it provides additional
+functionalities which are not available here. For more details about how to use DeepEcho
+whithin SDV, please visit the corresponding User Guide:
+
+* [SDV TimeSeries User Guide](https://sdv.dev/SDV/user_guides/timeseries/par.html)
+
+## Standalone usage
+
+**DeepEcho** can also be used as a standalone library.
+
+In this short quickstart, we show how to learn a mixed-type multivariate time series
+dataset and then generate synthetic data that resembles it.
+
+We will start by loading the data and preparing the instance of our model.
+
+```python3
+from deepecho import PARModel
+from deepecho.demo import load_demo
+
+# Load demo data
+data = load_demo()
+
+# Define data types for all the columns
+data_types = {
+    'region': 'categorical',
+    'day_of_week': 'categorical',
+    'total_sales': 'continuous',
+    'nb_customers': 'count',
+}
+
+model = PARModel(cuda=False)
+```
+
+If we want to use different settings for our model, like increasing the number
+of epochs or enabling CUDA, we can pass the arguments when creating the model:
+
+```python  # keep this as python (without the 3) to avoid using it in test-readme
+model = PARModel(epochs=1024, cuda=True)
+```
+
+Notice that for smaller datasets like the one used on this demo, CUDA usage introduces
+more overhead than the gains it obtains from parallelization, so the process in this
+case is more efficient without CUDA, even if it is available.
+
+Once we have created our instance, we are ready to learn the data and generate
+new synthetic data that resembles it:
+
+```python3
+# Learn a model from the data
+model.fit(
+    data=data,
+    entity_columns=['store_id'],
+    context_columns=['region'],
+    data_types=data_types,
+    sequence_index='date'
+)
+
+# Sample new data
+model.sample(num_entities=5)
+```
+
+The output will be a table with synthetic time series data with the same properties to
+the demo data that we used as input.
+
+# What's next?
+
+For more details about **DeepEcho** and all its possibilities and features, please check and
+run the [tutorials](tutorials).
+
+If you want to see how we evaluate the performance and quality of our models, please have a
+look at the [SDGym Benchmarking framework](https://github.com/sdv-dev/SDGym).
+
+Also, please feel welcome to visit [our contributing guide](CONTRIBUTING.rst) in order to help
+us developing new features or cool ideas!
+
+---
+
+
+<div align="center">
+<a href="https://datacebo.com"><img align="center" width=40% src="https://github.com/sdv-dev/SDV/blob/stable/docs/images/DataCebo.png"></img></a>
+</div>
+<br/>
+<br/>
+
+[The Synthetic Data Vault Project](https://sdv.dev) was first created at MIT's [Data to AI Lab](
+https://dai.lids.mit.edu/) in 2016. After 4 years of research and traction with enterprise, we
+created [DataCebo](https://datacebo.com) in 2020 with the goal of growing the project.
+Today, DataCebo is the proud developer of SDV, the largest ecosystem for
+synthetic data generation & evaluation. It is home to multiple libraries that support synthetic
+data, including:
+
+* 🔄 Data discovery & transformation. Reverse the transforms to reproduce realistic data.
+* 🧠 Multiple machine learning models -- ranging from Copulas to Deep Learning -- to create tabular,
+  multi table and time series data.
+* 📊 Measuring quality and privacy of synthetic data, and comparing different synthetic data
+  generation models.
+
+[Get started using the SDV package](https://sdv.dev/SDV/getting_started/install.html) -- a fully
+integrated solution and your one-stop shop for synthetic data. Or, use the standalone libraries
+for specific needs.

{deepecho-0.5.0.dev0 → deepecho-0.6.0.dev0}/deepecho/__init__.py

@@ -2,7 +2,7 @@
 
 __author__ = 'DataCebo, Inc.'
 __email__ = 'info@sdv.dev'
-__version__ = '0.5.0.dev0'
+__version__ = '0.6.0.dev0'
 __path__ = __import__('pkgutil').extend_path(__path__, __name__)
 
 from deepecho.demo import load_demo

{deepecho-0.5.0.dev0 → deepecho-0.6.0.dev0}/deepecho/models/base.py

@@ -6,7 +6,7 @@ import tqdm
 from deepecho.sequences import assemble_sequences
 
 
-class DeepEcho():
+class DeepEcho:
     """The base class for DeepEcho models."""
 
     _verbose = True
@@ -28,7 +28,13 @@ class DeepEcho():
             data_types:
                 See `fit`.
         """
-        dtypes = set(['continuous', 'categorical', 'ordinal', 'count', 'datetime'])
+        dtypes = set([
+            'continuous',
+            'categorical',
+            'ordinal',
+            'count',
+            'datetime',
+        ])
         assert all(dtype in dtypes for dtype in context_types)
         assert all(dtype in dtypes for dtype in data_types)
 
@@ -99,8 +105,15 @@ class DeepEcho():
 
         return dtypes_list
 
-    def fit(self, data, entity_columns=None, context_columns=None,
-            data_types=None, segment_size=None, sequence_index=None):
+    def fit(
+        self,
+        data,
+        entity_columns=None,
+        context_columns=None,
+        data_types=None,
+        segment_size=None,
+        sequence_index=None,
+    ):
         """Fit the model to a dataframe containing time series data.
 
         Args:
@@ -135,8 +148,7 @@ class DeepEcho():
         if segment_size is not None and not isinstance(segment_size, int):
             if sequence_index is None:
                 raise TypeError(
-                    '`segment_size` must be of type `int` if '
-                    'no `sequence_index` is given.'
+                    '`segment_size` must be of type `int` if ' 'no `sequence_index` is given.'
                 )
             if data[sequence_index].dtype.kind != 'M':
                 raise TypeError(
@@ -161,7 +173,12 @@ class DeepEcho():
         data_types = self._get_data_types(data, data_types, self._data_columns)
         context_types = self._get_data_types(data, data_types, self._context_columns)
         sequences = assemble_sequences(
-            data, self._entity_columns, self._context_columns, segment_size, sequence_index)
+            data,
+            self._entity_columns,
+            self._context_columns,
+            segment_size,
+            sequence_index,
+        )
 
         # Validate and fit
         self._validate(sequences, context_types, data_types)
@@ -242,7 +259,7 @@ class DeepEcho():
             # Reformat as a DataFrame
             group = pd.DataFrame(
                 dict(zip(self._data_columns, sequence)),
-                columns=self._data_columns
+                columns=self._data_columns,
             )
             group[self._entity_columns] = entity_values
             for column, value in zip(self._context_columns, context_values):

{deepecho-0.5.0.dev0 → deepecho-0.6.0.dev0}/deepecho/models/basic_gan.py

@@ -13,10 +13,13 @@ LOGGER = logging.getLogger(__name__)
 
 
 def _expand_context(data, context):
-    return torch.cat([
-        data,
-        context.unsqueeze(0).expand(data.shape[0], context.shape[0], context.shape[1])
-    ], dim=2)
+    return torch.cat(
+        [
+            data,
+            context.unsqueeze(0).expand(data.shape[0], context.shape[0], context.shape[1]),
+        ],
+        dim=2,
+    )
 
 
 class BasicGenerator(torch.nn.Module):
@@ -65,7 +68,7 @@ class BasicGenerator(torch.nn.Module):
         """
         latent = torch.randn(
             size=(sequence_length, context.size(0), self.latent_size),
-            device=self.device
+            device=self.device,
         )
         latent = _expand_context(latent, context)
 
@@ -150,8 +153,16 @@ class BasicGANModel(DeepEcho):
     _model_data_size = None
     _generator = None
 
-    def __init__(self, epochs=1024, latent_size=32, hidden_size=16,
-                 gen_lr=1e-3, dis_lr=1e-3, cuda=True, verbose=True):
+    def __init__(
+        self,
+        epochs=1024,
+        latent_size=32,
+        hidden_size=16,
+        gen_lr=1e-3,
+        dis_lr=1e-3,
+        cuda=True,
+        verbose=True,
+    ):
         self._epochs = epochs
         self._gen_lr = gen_lr
         self._dis_lr = dis_lr
@@ -211,7 +222,7 @@ class BasicGANModel(DeepEcho):
                     'type': column_type,
                     'min': np.min(values),
                     'max': np.max(values),
-                    'indices': (dimensions, dimensions + 1)
+                    'indices': (dimensions, dimensions + 1),
                 }
                 dimensions += 2
 
@@ -221,10 +232,7 @@ class BasicGANModel(DeepEcho):
                     indices[value] = dimensions
                     dimensions += 1
 
-                mapping[column] = {
-                    'type': column_type,
-                    'indices': indices
-                }
+                mapping[column] = {'type': column_type, 'indices': indices}
 
             else:
                 raise ValueError(f'Unsupported type: {column_type}')
@@ -317,7 +325,7 @@ class BasicGANModel(DeepEcho):
             self._one_hot_encode(tensor, value, properties)
 
         else:
-            raise ValueError()   # Theoretically unreachable
+            raise ValueError()  # Theoretically unreachable
 
     def _data_to_tensor(self, data):
         """Convert the input data to the corresponding tensor.
@@ -370,7 +378,7 @@ class BasicGANModel(DeepEcho):
                 elif column_type in ('categorical', 'ordinal'):
                     value = self._one_hot_decode(tensor, row, properties)
                 else:
-                    raise ValueError()   # Theoretically unreachable
+                    raise ValueError()  # Theoretically unreachable
 
                 column_data.append(value)
 
@@ -412,7 +420,7 @@ class BasicGANModel(DeepEcho):
             end_flag = sequence[:, self._data_size]
             if (end_flag == 1.0).any():
                 cut_idx = end_flag.detach().cpu().numpy().argmax()
-                sequence[cut_idx + 1:] = 0.0
+                sequence[cut_idx + 1 :] = 0.0
 
     def _generate(self, context, sequence_length=None):
         generated = self._generator(

{deepecho-0.5.0.dev0 → deepecho-0.6.0.dev0}/deepecho/models/par.py

@@ -27,10 +27,13 @@ class PARNet(torch.nn.Module):
         if isinstance(x, torch.nn.utils.rnn.PackedSequence):
             x, lengths = torch.nn.utils.rnn.pad_packed_sequence(x)
             if self.context_size:
-                x = torch.cat([
-                    x,
-                    c.unsqueeze(0).expand(x.shape[0], c.shape[0], c.shape[1])
-                ], dim=2)
+                x = torch.cat(
+                    [
+                        x,
+                        c.unsqueeze(0).expand(x.shape[0], c.shape[0], c.shape[1]),
+                    ],
+                    dim=2,
+                )
 
             x = self.down(x)
             x = torch.nn.utils.rnn.pack_padded_sequence(x, lengths, enforce_sorted=False)
@@ -41,10 +44,13 @@ class PARNet(torch.nn.Module):
 
         else:
             if self.context_size:
-                x = torch.cat([
-                    x,
-                    c.unsqueeze(0).expand(x.shape[0], c.shape[0], c.shape[1])
-                ], dim=2)
+                x = torch.cat(
+                    [
+                        x,
+                        c.unsqueeze(0).expand(x.shape[0], c.shape[0], c.shape[1]),
+                    ],
+                    dim=2,
+                )
 
             x = self.down(x)
             x, _ = self.rnn(x)
@@ -126,7 +132,7 @@ class PARModel(DeepEcho):
                     'mu': np.nanmean(x[i]),
                     'std': np.nanstd(x[i]),
                     'nulls': pd.isnull(x[i]).any(),
-                    'indices': (idx, idx + 1, idx + 2)
+                    'indices': (idx, idx + 1, idx + 2),
                 }
                 idx += 3
 
@@ -136,15 +142,12 @@ class PARModel(DeepEcho):
                     'min': np.nanmin(x[i]),
                     'range': np.nanmax(x[i]) - np.nanmin(x[i]),
                     'nulls': pd.isnull(x[i]).any(),
-                    'indices': (idx, idx + 1, idx + 2)
+                    'indices': (idx, idx + 1, idx + 2),
                 }
                 idx += 3
 
             elif t == 'categorical' or t == 'ordinal':
-                idx_map[i] = {
-                    'type': t,
-                    'indices': {}
-                }
+                idx_map[i] = {'type': t, 'indices': {}}
                 idx += 1
                 for v in set(x[i]):
                     if pd.isnull(v):
@@ -186,8 +189,8 @@ class PARModel(DeepEcho):
             'indices': {
                 '<START>': self._data_dims,
                 '<END>': self._data_dims + 1,
-                '<BODY>': self._data_dims + 2
-            }
+                '<BODY>': self._data_dims + 2,
+            },
         }
         self._data_dims += 3
 
@@ -225,7 +228,10 @@ class PARModel(DeepEcho):
                     x[p_idx] = 0.0
                     x[missing_idx] = 1.0 if pd.isnull(data[key][i]) else 0.0
 
-                elif props['type'] in ['categorical', 'ordinal']:   # categorical
+                elif props['type'] in [
+                    'categorical',
+                    'ordinal',
+                ]:  # categorical
                     value = data[key][i]
                     if pd.isnull(value):
                         value = None
@@ -250,15 +256,21 @@ class PARModel(DeepEcho):
         for key, props in self._ctx_map.items():
             if props['type'] in ['continuous', 'datetime']:
                 mu_idx, sigma_idx, missing_idx = props['indices']
-                x[mu_idx] = 0.0 if (pd.isnull(context[key]) or props['std'] == 0) else (
-                    context[key] - props['mu']) / props['std']
+                x[mu_idx] = (
+                    0.0
+                    if (pd.isnull(context[key]) or props['std'] == 0)
+                    else (context[key] - props['mu']) / props['std']
+                )
                 x[sigma_idx] = 0.0
                 x[missing_idx] = 1.0 if pd.isnull(context[key]) else 0.0
 
             elif props['type'] in ['count']:
                 r_idx, p_idx, missing_idx = props['indices']
-                x[r_idx] = 0.0 if (pd.isnull(context[key]) or props['range'] == 0) else (
-                    context[key] - props['min']) / props['range']
+                x[r_idx] = (
+                    0.0
+                    if (pd.isnull(context[key]) or props['range'] == 0)
+                    else (context[key] - props['min']) / props['range']
+                )
                 x[p_idx] = 0.0
                 x[missing_idx] = 1.0 if pd.isnull(context[key]) else 0.0
 
@@ -341,12 +353,13 @@ class PARModel(DeepEcho):
 
             epoch_loss_df = pd.DataFrame({
                 'Epoch': [epoch],
-                'Loss': [loss.item()]
+                'Loss': [loss.item()],
             })
             if not self.loss_values.empty:
-                self.loss_values = pd.concat(
-                    [self.loss_values, epoch_loss_df]
-                ).reset_index(drop=True)
+                self.loss_values = pd.concat([
+                    self.loss_values,
+                    epoch_loss_df,
+                ]).reset_index(drop=True)
             else:
                 self.loss_values = epoch_loss_df
 
@@ -375,7 +388,7 @@ class PARModel(DeepEcho):
                 This list contains the length of each sequence.
         """
         log_likelihood = 0.0
-        _, batch_size, input_size = X_padded.shape
+        _, batch_size, _input_size = X_padded.shape
 
         for key, props in self._data_map.items():
             if props['type'] in ['continuous', 'timestamp']:
@@ -386,14 +399,16 @@ class PARModel(DeepEcho):
 
                 for i in range(batch_size):
                     dist = torch.distributions.normal.Normal(
-                        mu[:seq_len[i], i], sigma[:seq_len[i], i])
-                    log_likelihood += torch.sum(dist.log_prob(X_padded[-seq_len[i]:, i, mu_idx]))
+                        mu[: seq_len[i], i], sigma[: seq_len[i], i]
+                    )
+                    log_likelihood += torch.sum(dist.log_prob(X_padded[-seq_len[i] :, i, mu_idx]))
 
-                    p_true = X_padded[:seq_len[i], i, missing_idx]
-                    p_pred = missing[:seq_len[i], i]
+                    p_true = X_padded[: seq_len[i], i, missing_idx]
+                    p_pred = missing[: seq_len[i], i]
                     log_likelihood += torch.sum(p_true * p_pred)
-                    log_likelihood += torch.sum((1.0 - p_true) * torch.log(
-                        1.0 - torch.exp(p_pred)))
+                    log_likelihood += torch.sum(
+                        (1.0 - p_true) * torch.log(1.0 - torch.exp(p_pred))
+                    )
 
             elif props['type'] in ['count']:
                 r_idx, p_idx, missing_idx = props['indices']
@@ -404,22 +419,26 @@ class PARModel(DeepEcho):
 
                 for i in range(batch_size):
                     dist = torch.distributions.negative_binomial.NegativeBinomial(
-                        r[:seq_len[i], i], p[:seq_len[i], i], validate_args=False)
-                    log_likelihood += torch.sum(dist.log_prob(x[:seq_len[i], i]))
-
-                    p_true = X_padded[:seq_len[i], i, missing_idx]
-                    p_pred = missing[:seq_len[i], i]
+                        r[: seq_len[i], i],
+                        p[: seq_len[i], i],
+                        validate_args=False,
+                    )
+                    log_likelihood += torch.sum(dist.log_prob(x[: seq_len[i], i]))
+
+                    p_true = X_padded[: seq_len[i], i, missing_idx]
+                    p_pred = missing[: seq_len[i], i]
                     log_likelihood += torch.sum(p_true * p_pred)
-                    log_likelihood += torch.sum((1.0 - p_true) * torch.log(
-                        1.0 - torch.exp(p_pred)))
+                    log_likelihood += torch.sum(
+                        (1.0 - p_true) * torch.log(1.0 - torch.exp(p_pred))
+                    )
 
             elif props['type'] in ['categorical', 'ordinal']:
                 idx = list(props['indices'].values())
                 log_softmax = torch.nn.functional.log_softmax(Y_padded[:, :, idx], dim=2)
 
                 for i in range(batch_size):
-                    target = X_padded[:seq_len[i], i, idx]
-                    predicted = log_softmax[:seq_len[i], i]
+                    target = X_padded[: seq_len[i], i, idx]
+                    predicted = log_softmax[: seq_len[i], i]
                     target = torch.argmax(target, dim=1).unsqueeze(dim=1)
                     log_likelihood += torch.sum(predicted.gather(dim=1, index=target))
 
@@ -443,14 +462,14 @@ class PARModel(DeepEcho):
             data[key] = []
             for i in range(seq_len):
                 if props['type'] in ['continuous', 'datetime']:
-                    mu_idx, sigma_idx, missing_idx = props['indices']
+                    mu_idx, _sigma_idx, missing_idx = props['indices']
                     if (x[i, 0, missing_idx] > 0) and props['nulls']:
                         data[key].append(None)
                     else:
                         data[key].append(x[i, 0, mu_idx].item() * props['std'] + props['mu'])
 
                 elif props['type'] in ['count']:
-                    r_idx, p_idx, missing_idx = props['indices']
+                    r_idx, _p_idx, missing_idx = props['indices']
                     if x[i, 0, missing_idx] > 0 and props['nulls']:
                         data[key].append(None)
                     else:
@@ -473,7 +492,7 @@ class PARModel(DeepEcho):
 
     def _sample_state(self, x):
         log_likelihood = 0.0
-        seq_len, batch_size, input_size = x.shape
+        seq_len, batch_size, _input_size = x.shape
         assert seq_len == 1 and batch_size == 1
 
         for key, props in self._data_map.items():