torch-mdct 0.1.0__tar.gz

torch_mdct-0.1.0/.gitignore

@@ -0,0 +1,160 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

torch_mdct-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Kinyugo Maina
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

torch_mdct-0.1.0/PKG-INFO

@@ -0,0 +1,56 @@
+Metadata-Version: 2.3
+Name: torch_mdct
+Version: 0.1.0
+Summary: A PyTorch implementation of the Modified Discrete Cosine Transform (MDCT) and its inverse for audio processing.
+Project-URL: Homepage, https://github.com/Kinyugo/torch_mdct
+Project-URL: Issues, https://github.com/Kinyugo/torch_mdct/issues
+Author-email: Kinyugo Maina <kinyugomaina@gmail.com>
+License-File: LICENSE
+Keywords: audio,imdct,mdct,pytorch,signal processing
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.8
+Requires-Dist: torch
+Provides-Extra: notebook
+Requires-Dist: ipython; extra == 'notebook'
+Requires-Dist: matplotlib; extra == 'notebook'
+Description-Content-Type: text/markdown
+
+# torch_mdct
+
+PyTorch implementation of Modified Discrete Cosine Transform and Inverse Modified Discrete Cosine Transform.
+
+## Installation 
+
+```bash
+pip install torch_mdct
+```
+
+## Usage
+
+```python
+import torchaudio
+from torch_mdct import MDCT, IMDCT
+
+# Load a sample waveform 
+waveform, sample_rate = torchaudio.load("/path/to/audio.file")
+
+# Initialize the mdct and imdct transforms
+mdct = MDCT(win_length=2048)
+imdct = IMDCT(win_length=2048)
+
+# Transform waveform into mdct spectrogram
+spectrogram = mdct(waveform)
+
+# Transform spectrogram back to audio 
+reconst_waveform = imdct(spectrogram)
+
+# Compute the differences
+print(f"L1: {(waveform - reconst_waveform).abs().mean()}")
+```
+
+## References 
+[[1]](https://github.com/zafarrafii/Zaf-Python) Zaf-Python: Zafar's Audio Functions in **Python** for audio signal analysis.
+
+[[2]](https://github.com/nils-werner/mdct) MDCT: A fast MDCT implementation using SciPy and FFTs.

torch_mdct-0.1.0/README.md

@@ -0,0 +1,37 @@
+# torch_mdct
+
+PyTorch implementation of Modified Discrete Cosine Transform and Inverse Modified Discrete Cosine Transform.
+
+## Installation 
+
+```bash
+pip install torch_mdct
+```
+
+## Usage
+
+```python
+import torchaudio
+from torch_mdct import MDCT, IMDCT
+
+# Load a sample waveform 
+waveform, sample_rate = torchaudio.load("/path/to/audio.file")
+
+# Initialize the mdct and imdct transforms
+mdct = MDCT(win_length=2048)
+imdct = IMDCT(win_length=2048)
+
+# Transform waveform into mdct spectrogram
+spectrogram = mdct(waveform)
+
+# Transform spectrogram back to audio 
+reconst_waveform = imdct(spectrogram)
+
+# Compute the differences
+print(f"L1: {(waveform - reconst_waveform).abs().mean()}")
+```
+
+## References 
+[[1]](https://github.com/zafarrafii/Zaf-Python) Zaf-Python: Zafar's Audio Functions in **Python** for audio signal analysis.
+
+[[2]](https://github.com/nils-werner/mdct) MDCT: A fast MDCT implementation using SciPy and FFTs.

torch_mdct-0.1.0/notebooks/example_usage.ipynb

@@ -0,0 +1,263 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Example Usage"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Uncomment to install the torch_mdct package\n",
+    "# %pip install torch_mdct"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Comment out the following line if you have installed the torch_mdct package\n",
+    "import sys\n",
+    "\n",
+    "sys.path.append(\"../src\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torchaudio\n",
+    "from IPython.display import Audio\n",
+    "from matplotlib import pyplot as plt\n",
+    "\n",
+    "from torch_mdct import IMDCT, MDCT"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Utils"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def plot_waveform(waveform: torch.Tensor, sample_rate: int, title: str) -> None:\n",
+    "    channels, n_frames = waveform.shape\n",
+    "\n",
+    "    skip = int(n_frames / (0.01 * n_frames))\n",
+    "    waveform = waveform[..., 0:-1:skip]\n",
+    "\n",
+    "    n_frames = waveform.shape[-1]\n",
+    "    time_axis = torch.linspace(0, n_frames / (sample_rate / skip), steps=n_frames)\n",
+    "\n",
+    "    fig, axes = plt.subplots(2, max(channels // 2, 1), constrained_layout=True)\n",
+    "    axes = axes.flatten()\n",
+    "\n",
+    "    for c in range(channels):\n",
+    "        axes[c].plot(time_axis, waveform[c], linewidth=1)\n",
+    "        axes[c].grid(True)\n",
+    "\n",
+    "        if channels > 1:\n",
+    "            axes[c].set_ylabel(f\"Channel {c}\")\n",
+    "\n",
+    "    fig.suptitle(title)\n",
+    "    plt.xlabel(\"Time (s)\")\n",
+    "    plt.show(block=False)\n",
+    "\n",
+    "\n",
+    "def plot_spectrogram(spectrogram: torch.Tensor, title: str) -> None:\n",
+    "    channels = spectrogram.shape[0]\n",
+    "\n",
+    "    fig, axes = plt.subplots(2, max(channels // 2, 1), constrained_layout=True)\n",
+    "    axes = axes.flatten()\n",
+    "\n",
+    "    for c in range(channels):\n",
+    "        im = axes[c].imshow(\n",
+    "            torch.log(spectrogram[c].abs() + 1e-5), origin=\"lower\", aspect=\"auto\"\n",
+    "        )\n",
+    "        fig.colorbar(im, ax=axes[c])\n",
+    "\n",
+    "        if channels > 1:\n",
+    "            axes[c].set_ylabel(f\"Channels {c}\")\n",
+    "\n",
+    "    fig.suptitle(title)\n",
+    "    plt.xlabel(\"Time\")\n",
+    "    plt.show(block=False)\n",
+    "\n",
+    "\n",
+    "def stats(x: torch.Tensor) -> str:\n",
+    "    return f\"Shape: {tuple(x.shape)} Min: {x.min():.2f} Max: {x.max():.2f} Mean: {x.mean():.2f} Std: {x.std():.2f}\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## DataLoading"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "waveform, sample_rate = torchaudio.load(\"audio_samples/sample.ogg\")\n",
+    "Audio(waveform, rate=sample_rate)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plot_waveform(waveform, sample_rate, f\"Waveform: \\n({stats(waveform)})\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Transforms"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mdct = MDCT(win_length=2048)\n",
+    "imdct = IMDCT(win_length=2048)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## MDCT Experiments"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "spectrogram = mdct(waveform)\n",
+    "plot_spectrogram(spectrogram, f\"Log Absolute Spectrogram: \\n({stats(spectrogram)})\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## IMDCT Experiments"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "reconst_waveform = imdct(spectrogram, length=waveform.shape[-1])\n",
+    "Audio(reconst_waveform, rate=sample_rate)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plot_waveform(\n",
+    "    reconst_waveform,\n",
+    "    sample_rate,\n",
+    "    f\"Reconstructed Waveform: \\n({stats(reconst_waveform)})\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Waveform Difference"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "waveform_diff = waveform - reconst_waveform\n",
+    "Audio(waveform_diff, rate=sample_rate)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plot_waveform(\n",
+    "    waveform_diff,\n",
+    "    sample_rate,\n",
+    "    f\"Waveform Difference: \\n({stats(waveform_diff)})\",\n",
+    ")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "torch",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

torch_mdct-0.1.0/pyproject.toml

@@ -0,0 +1,26 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "torch_mdct"
+version = "0.1.0"
+authors = [{name = "Kinyugo Maina", email = "kinyugomaina@gmail.com"}]
+description = "A PyTorch implementation of the Modified Discrete Cosine Transform (MDCT) and its inverse for audio processing."
+readme = "README.md"
+licence = {file = "LICENSE"}
+requires-python = ">=3.8"
+keywords = ["pytorch", "mdct", "imdct", "audio", "signal processing"]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+dependencies = ["torch"]
+
+[project.optional-dependencies]
+notebook = ["matplotlib", "ipython"]
+
+[project.urls]
+Homepage = "https://github.com/Kinyugo/torch_mdct"
+Issues = "https://github.com/Kinyugo/torch_mdct/issues"

torch_mdct-0.1.0/src/torch_mdct/__init__.py

@@ -0,0 +1,3 @@
+from .functional import imdct, mdct
+from .torch_mdct import IMDCT, MDCT
+from .windows import kaiser_bessel_derived, vorbis

torch_mdct-0.1.0/src/torch_mdct/functional.py

@@ -0,0 +1,145 @@
+import math
+
+import torch
+from torch.nn import functional as F
+
+
+def mdct(
+    waveform: torch.Tensor, window: torch.Tensor, center: bool = True
+) -> torch.Tensor:
+    """
+    Compute the Modified Discrete Cosine Transform (MDCT) of a waveform.
+
+    Parameters
+    ----------
+    waveform : torch.Tensor
+        Input waveform tensor of shape (..., n_samples).
+    window : torch.Tensor
+        Window tensor of shape (win_length,).
+    center : bool, optional
+        If True, pad the waveform on both sides with half the window length, by default True.
+
+    Returns
+    -------
+    torch.Tensor
+        MDCT spectrogram of the input waveform of shape (..., win_length // 2, n_frames).
+    """
+    n_samples = waveform.shape[-1]
+    win_length = window.shape[-1]
+    hop_length = win_length // 2
+
+    # Flatten the input tensor
+    shape = waveform.shape
+    waveform = waveform.flatten(end_dim=-2)
+
+    # Center the waveform by padding on both sides
+    n_frames = int(math.ceil(n_samples / hop_length)) + 1
+    if center:
+        waveform = F.pad(
+            waveform,
+            (hop_length, (n_frames + 1) * hop_length - n_samples),
+            mode="constant",
+        )
+
+    # Prepare pre&post processing arrays
+    pre_twiddle = torch.exp(
+        -1j
+        * torch.pi
+        / win_length
+        * torch.arange(0, win_length, device=waveform.device)
+    )
+    post_twiddle = torch.exp(
+        -1j
+        * torch.pi
+        / win_length
+        * (win_length / 2 + 1)
+        * torch.arange(0.5, win_length / 2 + 0.5, device=waveform.device)
+    )
+
+    # Convolve the waveform with the window and fft matrices
+    spectrogram = waveform.unfold(dimension=-1, size=win_length, step=hop_length)
+    spectrogram = torch.einsum("...kj,j->...kj", spectrogram, window)
+    spectrogram = torch.einsum("...kj,j->...kj", spectrogram, pre_twiddle)
+    spectrogram = torch.fft.fft(spectrogram, dim=-1)
+    spectrogram = torch.einsum(
+        "...jk,k->...kj", spectrogram[..., : win_length // 2], post_twiddle
+    )
+    spectrogram = torch.real(spectrogram)
+
+    # Unflatten the output
+    spectrogram = spectrogram.reshape(shape[:-1] + spectrogram.shape[-2:])
+
+    return spectrogram
+
+
+def imdct(
+    spectrogram: torch.Tensor, window: torch.Tensor, center: bool = True
+) -> torch.Tensor:
+    """
+    Compute the inverse Modified Discrete Cosine Transform (iMDCT) of a spectrogram.
+
+    Parameters
+    ----------
+    spectrogram : torch.Tensor
+        Input MDCT spectrogram tensor of shape (..., win_length // 2, n_frames).
+    window : torch.Tensor
+        Window tensor of shape (win_length,).
+    center : bool, optional
+        If True, remove the padding added during MDCT, by default True.
+
+    Returns
+    -------
+    torch.Tensor
+        Reconstructed waveform tensor of shape (..., n_samples).
+    """
+    win_length = window.shape[-1]
+    hop_length = win_length // 2
+    n_freqs, n_frames = spectrogram.shape[-2:]
+    n_samples = hop_length * (n_frames + 1)
+
+    # Flatten the input tensor
+    shape = spectrogram.shape
+    spectrogram = spectrogram.flatten(end_dim=-3)
+
+    # Prepare pre&post processing arrays
+    pre_twiddle = torch.exp(
+        -1j
+        * torch.pi
+        / (2 * n_freqs)
+        * (n_freqs + 1)
+        * torch.arange(n_freqs, device=spectrogram.device)
+    )
+    post_twiddle = (
+        torch.exp(
+            -1j
+            * torch.pi
+            / (2 * n_freqs)
+            * torch.arange(
+                0.5 + n_freqs / 2,
+                2 * n_freqs + n_freqs / 2 + 0.5,
+                device=spectrogram.device,
+            )
+        )
+        / n_freqs
+    )
+
+    # Apply fft and the window
+    spectrogram = torch.einsum("...jk,j->...jk", spectrogram, pre_twiddle)
+    spectrogram = torch.fft.fft(spectrogram, n=2 * n_freqs, dim=1)
+    spectrogram = torch.einsum("...jk,j->...jk", spectrogram, post_twiddle)
+    spectrogram = torch.real(spectrogram)
+    spectrogram = 2 * torch.einsum("...jk,j->...jk", spectrogram, window)
+
+    # Recover the waveform with the time-domain aliasing cancelling principle
+    waveform = F.fold(
+        spectrogram, (1, n_samples), kernel_size=(1, win_length), stride=(1, hop_length)
+    )
+
+    # Remove padding
+    if center:
+        waveform = waveform[..., hop_length:-hop_length]
+
+    # Unflatten the output
+    waveform = waveform.reshape((*shape[:-2], -1))
+
+    return waveform

torch_mdct-0.1.0/src/torch_mdct/torch_mdct.py

@@ -0,0 +1,128 @@
+from functools import partial
+from typing import Any, Dict, Optional, Union
+
+import torch
+from torch import nn
+
+from .functional import imdct, mdct
+from .windows import kaiser_bessel_derived, vorbis
+
+
+class MDCT(nn.Module):
+    """
+    Module to compute the Modified Discrete Cosine Transform (MDCT) of a waveform.
+
+    Parameters
+    ----------
+    win_length : int
+        The length of the window.
+    window_fn : callable, optional
+        A function to generate the window, by default vorbis.
+    window_kwargs : dict, optional
+        Additional keyword arguments to pass to the window function, by default {}.
+    center : bool, optional
+        If True, pad the waveform on both sides with half the window length, by default True.
+
+    Attributes
+    ----------
+    window : torch.Tensor
+        The window tensor.
+    mdct : callable
+        The MDCT function with the specified center parameter.
+
+    Methods
+    -------
+    forward(waveform: torch.Tensor) -> torch.Tensor
+        Compute the MDCT of the input waveform.
+    """
+
+    def __init__(
+        self,
+        win_length: int,
+        window_fn: Union[kaiser_bessel_derived, vorbis] = vorbis,
+        window_kwargs: Dict[str, Any] = {},
+        center: bool = True,
+    ) -> None:
+        super().__init__()
+
+        self.register_buffer("window", window_fn(win_length, **window_kwargs))
+        self.mdct = partial(mdct, center=center)
+
+    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the MDCT of the input waveform.
+
+        Parameters
+        ----------
+        waveform : torch.Tensor
+            Input waveform tensor of shape (..., n_samples).
+
+        Returns
+        -------
+        torch.Tensor
+            MDCT spectrogram of the input waveform of shape (..., win_length // 2, n_frames).
+        """
+        return self.mdct(waveform, self.window)
+
+
+class IMDCT(nn.Module):
+    """
+    Module to compute the inverse Modified Discrete Cosine Transform (iMDCT) of a spectrogram.
+
+    Parameters
+    ----------
+    win_length : int
+        The length of the window.
+    window_fn : callable, optional
+        A function to generate the window, by default vorbis.
+    window_kwargs : dict, optional
+        Additional keyword arguments to pass to the window function, by default {}.
+    center : bool, optional
+        If True, remove the padding added during MDCT, by default True.
+
+    Attributes
+    ----------
+    window : torch.Tensor
+        The window tensor.
+    imdct : callable
+        The iMDCT function with the specified center parameter.
+
+    Methods
+    -------
+    forward(spectrogram: torch.Tensor) -> torch.Tensor
+        Compute the iMDCT of the input spectrogram.
+    """
+
+    def __init__(
+        self,
+        win_length: int,
+        window_fn: Union[kaiser_bessel_derived, vorbis] = vorbis,
+        window_kwargs: Dict[str, Any] = {},
+        center: bool = True,
+    ) -> None:
+        super().__init__()
+
+        self.register_buffer("window", window_fn(win_length, **window_kwargs))
+        self.imdct = partial(imdct, center=center)
+
+    def forward(
+        self, spectrogram: torch.Tensor, length: Optional[int] = None
+    ) -> torch.Tensor:
+        """
+        Compute the iMDCT of the input spectrogram.
+
+        Parameters
+        ----------
+        spectrogram : torch.Tensor
+            Input MDCT spectrogram tensor of shape (..., win_length // 2, n_frames).
+
+        length : int, optional
+            The length of the output waveform, by default None.
+
+        Returns
+        -------
+        torch.Tensor
+            Reconstructed waveform tensor of shape (..., n_samples).
+        """
+        length = length or -1  # By default, return the full waveform
+        return self.imdct(spectrogram, self.window)[..., :length]

torch_mdct-0.1.0/src/torch_mdct/windows.py

@@ -0,0 +1,68 @@
+import torch
+
+
+def kaiser_bessel_derived(
+    win_length: int,
+    beta: float = 12.0,
+    *,
+    dtype: torch.dtype = None,
+    device: torch.device = None
+) -> torch.Tensor:
+    """
+    Generate a Kaiser-Bessel derived window.
+
+    Parameters
+    ----------
+    win_length : int
+        The length of the window.
+    beta : float, optional
+        The beta parameter for the Kaiser window, by default 12.0.
+    dtype : torch.dtype, optional
+        The desired data type of returned tensor, by default None.
+    device : torch.device, optional
+        The desired device of returned tensor, by default None.
+
+    Returns
+    -------
+    torch.Tensor
+        The generated Kaiser-Bessel derived window of shape (win_length,).
+    """
+
+    half_w_length = win_length // 2
+    kaiser_w = torch.kaiser_window(
+        half_w_length + 1, True, beta, dtype=dtype, device=device
+    )
+    kaiser_w_csum = torch.cumsum(kaiser_w, dim=-1)
+    half_w = torch.sqrt(kaiser_w_csum[:-1] / kaiser_w_csum[-1])
+    w = torch.cat((half_w, torch.flip(half_w, dims=(0,))), axis=0)
+
+    return w
+
+
+def vorbis(
+    win_length: int, *, dtype: torch.dtype = None, device: torch.device = None
+) -> torch.Tensor:
+    """
+    Generate a Vorbis window.
+
+    Parameters
+    ----------
+    win_length : int
+        The length of the window.
+    dtype : torch.dtype, optional
+        The desired data type of returned tensor, by default None.
+    device : torch.device, optional
+        The desired device of returned tensor, by default None.
+
+    Returns
+    -------
+    torch.Tensor
+        The generated Vorbis window of shape (win_length,).
+    """
+
+    arg = torch.arange(win_length, dtype=dtype, device=device) + 0.5
+    w = torch.sin(
+        torch.pi / 2.0 * torch.pow(torch.sin(torch.pi / win_length * arg), 2.0)
+    )
+
+    return w