update

.gitignore

@@ -0,0 +1,168 @@
+# 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/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# 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/
+
+toy_exp/
+Checkpoints*/
+Data/datasets/
+wandb/
+data/

Infer Demo.ipynb

@@ -0,0 +1,169 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import torch\n",
+    "\n",
+    "torch.cuda.is_available()\n",
+    "os.environ[\"WANDB_ENABLED\"] = \"false\"\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "device = torch.device(f\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "from models.Tiffusion import tiffusion\n",
+    "# from models.CSDI import tiffusion\n",
+    "\n",
+    "model = tiffusion.Tiffusion(\n",
+    "    seq_length=365,\n",
+    "    feature_size=3,\n",
+    "    n_layer_enc=6,\n",
+    "    n_layer_dec=4,\n",
+    "    d_model=128,\n",
+    "    timesteps=500,\n",
+    "    sampling_timesteps=200,\n",
+    "    loss_type='l1',\n",
+    "    beta_schedule='cosine',\n",
+    "    n_heads=8,\n",
+    "    mlp_hidden_times=4,\n",
+    "    attn_pd=0.0,\n",
+    "    resid_pd=0.0,\n",
+    "    kernel_size=1,\n",
+    "    padding_size=0,\n",
+    "    control_signal=[]\n",
+    ").to(device)\n",
+    "\n",
+    "model.load_state_dict(torch.load(\"./weight/checkpoint-10.pt\", map_location='cpu', weights_only=True)[\"model\"])\n",
+    "# model.load_state_dict(torch.load(\"../../../data/CSDI/ckpt_baseline_365/checkpoint-10.pt\", map_location='cpu', weights_only=True)[\"model\"])\n",
+    "\n",
+    "\n",
+    "coef = 1.0e-2\n",
+    "stepsize =  5.0e-2\n",
+    "sampling_steps = 100 # 这个可以调整 100-500都行 快慢和准度 tradeoff\n",
+    "seq_length = 365\n",
+    "feature_dim = 3\n",
+    "print(f\"seq_length: {seq_length}, feature_dim: {feature_dim}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Sampling"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "anchor_value = [\n",
+    "    # (time, feature_id, y-value, confidence)\n",
+    "    (0, 0, 0.04, 1.0),\n",
+    "    (2, 0, 0.58, 1.0),\n",
+    "    # (6, 0, 0.27, 0.5),\n",
+    "    # (10, 0, 0.04, 1.0),\n",
+    "    # (12, 0, 0.58, 0.001),\n",
+    "    # (16, 0, 0.27, 0.5),\n",
+    "    # (20, 0, 0.04, 1.0),\n",
+    "    # (22, 0, 0.58, 0.001),\n",
+    "    # (26, 0, 0.27, 0.5),\n",
+    "    # (30, 0, 0.04, 1.0),\n",
+    "    # (32, 0, 0.58, 0.001),\n",
+    "    # (36, 0, 0.27, 0.5),\n",
+    "    # (40, 0, 0.04, 1.0),\n",
+    "    # (42, 0, 0.58, 0.001),\n",
+    "    # (46, 0, 0.27, 0.5),\n",
+    "    # (50, 0, 0.04, 1.0),\n",
+    "    # (52, 0, 0.58, 0.001),\n",
+    "    # (56, 0, 0.27, 0.5),\n",
+    "    # (60, 0, 0.04, 1.0),\n",
+    "    # (62, 0, 0.58, 0.001),\n",
+    "    # (66, 0, 0.27, 0.5),\n",
+    "]\n",
+    "\n",
+    "observed_points = torch.zeros((seq_length, feature_dim)).to(device)\n",
+    "observed_mask = torch.zeros((seq_length, feature_dim)).to(device)\n",
+    "\n",
+    "for time, feature_id, y_value, confidence in anchor_value:\n",
+    "    observed_points[time, feature_id] = y_value\n",
+    "    observed_mask[time, feature_id] = confidence\n",
+    "\n",
+    "auc = -10\n",
+    "auc_weight = 10.0\n",
+    "with torch.no_grad():\n",
+    "    results = model.predict_weighted_points(\n",
+    "        observed_points, # (seq_length, feature_dim)\n",
+    "        observed_mask, # (seq_length, feature_dim)\n",
+    "        coef,  # fixed\n",
+    "        stepsize, # fixed\n",
+    "        sampling_steps, # fixed\n",
+    "        # model_control_signal=model_control_signal,\n",
+    "        gradient_control_signal={\n",
+    "            \"auc\": auc, \"auc_weight\": auc_weight,\n",
+    "        },\n",
+    "    )\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "results.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "plt.plot(results[:,0], label=\"Predicted Feature 0\")\n",
+    "for time, feature_id, y_value, confidence in anchor_value:\n",
+    "    plt.scatter(time, y_value, c='r')\n",
+    "plt.show()\n",
+    "plt.plot(results[:,1], label=\"Predicted Feature 1\")\n",
+    "plt.show()\n",
+    "plt.plot(results[:,2], label=\"Predicted Feature 2\")\n",
+    "plt.show()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "rag",
+   "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.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 XXX
+
+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.

README.md

@@ -1,13 +1,128 @@
----
-title: TSEditor
-emoji: 👁
-colorFrom: purple
-colorTo: pink
-sdk: gradio
-sdk_version: 5.13.1
-app_file: app.py
-pinned: false
-short_description: TSEditor Demo for ICML 2025
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Demo for TSEditor
+<!-- 
+[![](https://img.shields.io/github/stars/Y-debug-sys/Diffusion-TS.svg)](https://github.com/Y-debug-sys/Diffusion-TS/stargazers)
+[![](https://img.shields.io/github/forks/Y-debug-sys/Diffusion-TS.svg)](https://github.com/Y-debug-sys/Diffusion-TS/network) 
+[![](https://img.shields.io/badge/license-MIT-blue.svg)](https://github.com/Y-debug-sys/Diffusion-TS/blob/main/LICENSE) 
+<img src="https://img.shields.io/badge/python-3.8-blue">
+<img src="https://img.shields.io/badge/pytorch-2.0-orange">
+
+> **Abstract:** Denoising diffusion probabilistic models (DDPMs) are becoming the leading paradigm for generative models. It has recently shown breakthroughs in audio synthesis, time series imputation and forecasting. In this paper, we propose Diffusion-TS, a novel diffusion-based framework that generates multivariate time series samples of high quality by using an encoder-decoder transformer with disentangled temporal representations, in which the decomposition technique guides Diffusion-TS to capture the semantic meaning of time series while transformers mine detailed sequential information from the noisy model input. Different from existing diffusion-based approaches, we train the model to directly reconstruct the sample instead of the noise in each diffusion step, combining a Fourier-based loss term. Diffusion-TS is expected to generate time series satisfying both interpretablity and realness. In addition, it is shown that the proposed Diffusion-TS can be easily extended to conditional generation tasks, such as forecasting and imputation, without any model changes. This also motivates us to further explore the performance of Diffusion-TS under irregular settings. Finally, through qualitative and quantitative experiments, results show that Diffusion-TS achieves the state-of-the-art results on various realistic analyses of time series. 
+
+Diffusion-TS is a diffusion-based framework that generates general time series samples both conditionally and unconditionally. As shown in Figure 1, the framework contains two parts: a sequence encoder and an interpretable decoder which decomposes the time series into seasonal part and trend part. The trend part contains the polynomial regressor and extracted mean of each block output. For seasonal part, we reuse trigonometric representations based on Fourier series. Regarding training, sampling and more details, please refer to [our paper](https://openreview.net/pdf?id=4h1apFjO99) in ICLR 2024. 
+
+<p align="center">
+  <img src="figures/fig1.jpg" alt="">
+  <br>
+  <b>Figure 1</b>: Overall Architecture of Diffusion-TS.
+</p>
+
+
+## Dataset Preparation
+
+All the four real-world datasets (Stocks, ETTh1, Energy and fMRI) can be obtained from [Google Drive](https://drive.google.com/file/d/11DI22zKWtHjXMnNGPWNUbyGz-JiEtZy6/view?usp=sharing). Please download **dataset.zip**, then unzip and copy it to the folder `./Data` in our repository.
+
+
+## Running the Code
+
+ The code requires conda3 (or miniconda3), and one CUDA capable GPU. The instructions below guide you regarding running the codes in this repository. 
+
+### Environment & Libraries
+
+The full libraries list is provided as a `requirements.txt` in this repo. Please create a virtual environment with `conda` or `venv` and run
+
+~~~bash
+(myenv) $ pip install -r requirements.txt
+~~~
+
+### Training & Sampling
+
+For training, you can reproduce the experimental results of all benchmarks by runing
+
+~~~bash
+(myenv) $ python main.py --name {name} --config_file {config.yaml} --gpu 0 --train
+~~~
+
+**Note:** We also provided the corresponding `.yml` files (only stocks, sines, mujoco, etth, energy and fmri) under the folder `./Config` where all possible option can be altered. You may need to change some parameters in the model for different scenarios. For example, we use the whole data to train model for unconditional evaluation, then *training_ratio* is set to 1 by default. As for conditional generation, we need to divide data set thus it should be changed to a value < 1. 
+
+While training, the script will save check points to the *results* folder after a fixed number of epochs. Once trained, please use the saved model for sampling by running
+
+#### Unconstrained
+```bash
+(myenv) $ python main.py --name {name} --config_file {config.yaml} --gpu 0 --sample 0 --milestone {checkpoint_number}
+```
+
+#### Imputation
+```bash
+(myenv) $ python main.py --name {name} --config_file {config.yaml} --gpu 0 --sample 1 --milestone {checkpoint_number} --mode infill --missing_ratio {missing_ratio}
+```
+
+#### Forecasting
+```bash
+(myenv) $ python main.py --name {dataset_name} --config_file {config.yaml} --gpu 0 --sample 1 --milestone {checkpoint_number} --mode predict --pred_len {pred_len}
+```
+
+
+## Visualization and Evaluation
+
+After sampling, synthetic data and orginal data are stored in `.npy` file format under the *output* folder, which can be directly read to calculate quantitative metrics such as discriminative, predictive, correlational and context-FID score. You can also reproduce the visualization results using t-SNE or kernel plotting, and all of these evaluational codes can be found in the folder `./Utils`. Please refer to `.ipynb` tutorial files in this repo for more detailed implementations.
+
+**Note:** All the metrics can be found in the `./Experiments` folder. Additionally, by default, for datasets other than the Sine dataset (because it do not need normalization), their normalized forms are saved in `{...}_norm_truth.npy`. Therefore, when you run the Jupternotebook for dataset other than Sine, just uncomment and rewrite the corresponding code written at the beginning.
+
+### Main Results
+
+#### Standard TS Generation
+<p align="center">
+  <b>Table 1</b>: Results of 24-length Time-series Generation.
+  <br>
+  <img src="figures/fig2.jpg" alt="">
+</p>
+
+#### Long-term TS Generation
+<p align="center">
+  <b>Table 2</b>: Results of Long-term Time-series Generation.
+  <br>
+  <img src="figures/fig3.jpg" alt="">
+</p>
+
+#### Conditional TS Generation
+<p align="center">
+  <img src="figures/fig4.jpg" alt="">
+  <br>
+  <b>Figure 2</b>: Visualizations of Time-series Imputation and Forecasting.
+</p>
+
+
+## Authors
+
+* Paper Authors : Xinyu Yuan, Yan Qiao
+
+* Code Author : Xinyu Yuan
+
+* Contact : yxy5315@gmail.com
+ -->
+
+## Citation
+If you find this repo useful, please cite our paper via
+```bibtex
+```
+
+
+## Acknowledgement
+
+We appreciate the following github repos a lot for their valuable code base:
+
+https://github.com/Y-debug-sys/Diffusion-TS
+
+https://github.com/lucidrains/denoising-diffusion-pytorch
+
+https://github.com/cientgu/VQ-Diffusion
+
+https://github.com/XiangLi1999/Diffusion-LM
+
+https://github.com/philipperemy/n-beats
+
+https://github.com/salesforce/ETSformer
+
+https://github.com/ermongroup/CSDI
+
+https://github.com/jsyoon0823/TimeGAN

app.py

@@ -0,0 +1,1121 @@
+import gradio as gr
+import pandas as pd
+import plotly.graph_objects as go
+import numpy as np
+import torch
+from typing import Dict, List, Tuple
+import re
+from typing import Callable, Union, Dict
+
+
+class TimeSeriesEditor:
+    def __init__(self, seq_length: int, feature_dim: int, trainer):
+        # Existing initialization
+        self.seq_length = seq_length
+        self.feature_dim = feature_dim
+        self.trainer = trainer
+        self.coef = None
+        self.stepsize = None
+        self.sampling_steps = None
+        self.feature_names = ["revenue", "download", "daily active user"]# * 20
+        # self.feature_names = [f"Feature {i}" for i in range(self.feature_dim)]
+
+        # Store the latest model output
+        self.latest_sample = None
+        self.latest_observed_points = None
+        self.latest_observed_mask = None
+        self.latest_gradient_control_signal = None
+        self.latest_model_control_signal = None
+        # self.latest_metrics 
+        # Define scales for each feature
+        self.feature_scales = {
+            0: 1000000,  # Revenue: $1M per 0.1
+            1: 100000,   # Download: 100K downloads per 0.1
+            2: 10000     # AU: 10K active users per 0.1
+        }
+        self.feature_units = {
+            0: "$",          # Revenue
+            1: "downloads",  # Download
+            2: "users"       # AU
+        }
+        self.show_normalized = True
+
+        # Add frequency band multipliers
+        self.freq_bands = np.ones(5)  # 5 frequency bands, initially all set to 1.0
+        self.function_parser = FunctionParser()
+        self.trending_controls = []  # Store trending controls
+        
+    def format_value(self, value: float, feature_idx: int) -> str:
+        """Format value with appropriate units and notation"""
+        if self.show_normalized:
+            return f"{value:.4f}"
+        else:
+            if feature_idx == 0:  # Revenue
+                return f"{self.feature_units[feature_idx]}{value:,.2f}"
+            else:  # Downloads and AU
+                return f"{value:,.0f} {self.feature_units[feature_idx]}"
+
+    def create_plot(self, sample: np.ndarray, observed_points: torch.Tensor,
+                    observed_mask: torch.Tensor,
+                gradient_control_signal: Dict, metrics: Dict) -> List[go.Figure]:
+        figures = []
+        # Get weights from model_control_signal (will be all 1s if not provided)
+        weights = observed_mask
+
+        for feat_idx in range(self.feature_dim):
+            fig = go.Figure()
+            
+            # Scale values if needed
+            scale_factor = self.feature_scales[feat_idx] * 10 if not self.show_normalized else 1
+            
+            # Plot predicted line
+            predicted_values = sample[:, feat_idx] * scale_factor
+            fig.add_trace(go.Scatter(
+                x=np.arange(self.seq_length),
+                y=predicted_values,
+                mode='lines',
+                name='Predicted',
+                line=dict(color='green', width=2),
+                showlegend=True
+            ))
+            
+            # Calculate and plot confidence bands based on weights
+            # Lower weights = larger uncertainty bands
+            mask = observed_points[:, feat_idx] > 0
+            ox = np.arange(0, self.seq_length)[mask]
+            oy = observed_points[mask, feat_idx].numpy() * scale_factor
+            weights_masked = 1 - weights[mask, feat_idx].numpy()
+
+            # Calculate error bars - inverse relationship with weight
+            # Weight of 1.0 gives minimal uncertainty (0.02)
+            # Weight of 0.1 gives larger uncertainty (0.2)
+            # error_y = 0.02 / weights_masked
+            error_y = weights_masked / 5
+
+            # Plot observed points with error bars - changed symbol to 'cross'
+            fig.add_trace(go.Scatter(
+                x=ox,
+                y=oy,
+                mode='markers',
+                name='Observed',
+                marker=dict(
+                    # special red
+                    color='rgba(255, 0, 0, 0.5)',
+                    # size=10,
+                    symbol='x',  # Changed from 'circle' to 'x' for cross symbol
+                ),
+                error_y=dict(
+                    type='data',
+                    array=error_y * scale_factor,
+                    visible=True,
+                    thickness=0.5,
+                    width=2,
+                    color='blue'
+                ),
+                showlegend=True
+            ))
+            
+            # Add shaded confidence bands around the predicted line
+            # This shows the general uncertainty in the prediction
+            uncertainty = 0.05  # Base uncertainty level
+            upper_bound = predicted_values + uncertainty * scale_factor
+            lower_bound = predicted_values - uncertainty * scale_factor
+            
+            fig.add_trace(go.Scatter(
+                x=np.concatenate([np.arange(self.seq_length), np.arange(self.seq_length)[::-1]]),
+                y=np.concatenate([upper_bound, lower_bound[::-1]]),
+                # fill='toself',
+                # fillcolor='rgba(0,100,0,0.1)',
+                line=dict(color='rgba(255,255,255,0)'),
+                name='Prediction Interval',
+                showlegend=True
+            ))
+            
+            # Add vertical lines for peak points
+            if gradient_control_signal.get("peak_points"):
+                for peak_point in gradient_control_signal["peak_points"]:
+                    fig.add_vline(x=peak_point, line_dash="dash", line_color="red")
+            
+            # Add metrics annotations
+            total_value = np.sum(sample[:, feat_idx]) * (self.feature_scales[feat_idx] * 10 if not self.show_normalized else 1)
+            annotations = [dict(
+                x=0.02,
+                y=1.1,
+                xref="paper",
+                yref="paper",
+                text=f"Total {self.feature_names[feat_idx]}: {self.format_value(total_value, feat_idx)}",
+                showarrow=False
+            )]
+            
+            # Update y-axis title based on feature and scaling
+            if self.show_normalized:
+                y_title = f'{self.feature_names[feat_idx]} (Normalized)'
+            else:
+                unit = self.feature_units[feat_idx]
+                y_title = f'{self.feature_names[feat_idx]} ({unit})'
+            
+            # Create a more informative legend for uncertainty
+            legend_text = (
+                "Prediction with Confidence Bands<br>"
+                "• Blue points: Observed values with uncertainty<br>"
+                "• Green line: Predicted values<br>"
+                # "• Shaded area: Prediction uncertainty<br>"
+                "• Error bars: Observation uncertainty (larger = lower weight)"
+            )
+            
+            fig.update_layout(
+                title=dict(
+                    text=f'Feature: {self.feature_names[feat_idx]}',
+                    x=0.5,
+                    y=0.95
+                ),
+                xaxis_title='Time',
+                yaxis_title=y_title,
+                height=400,
+                showlegend=True,
+                dragmode='select',
+                annotations=[
+                    *annotations,
+                    # dict(
+                    #     x=1.15,
+                    #     y=0.5,
+                    #     xref="paper",
+                    #     yref="paper",
+                    #     text=legend_text,
+                    #     showarrow=False,
+                    #     align="left",
+                    #     bordercolor="black",
+                    #     borderwidth=1,
+                    #     borderpad=4,
+                    #     bgcolor="white",
+                    # )
+                ],
+                margin=dict(r=200)  # Add right margin for legend
+            )
+            
+            figures.append(fig)
+
+        return figures
+
+    def update_scaling(self,
+                      revenue_scale: float,
+                      download_scale: float,
+                      au_scale: float,
+                      show_normalized: bool) -> Tuple[List[go.Figure], Dict]:
+        """Update the scaling parameters and redraw plots"""
+        if self.latest_sample is None:
+            return [], {}
+
+        # Update scales
+        self.feature_scales = {
+            0: revenue_scale,
+            1: download_scale,
+            2: au_scale
+        }
+        self.show_normalized = show_normalized
+        
+        # Calculate metrics
+        metrics = {
+            'show_normalized': self.show_normalized
+        }
+        for feat_idx in range(self.feature_dim):
+            total = np.sum(self.latest_sample[:, feat_idx]) * (self.feature_scales[feat_idx] * 10 if not self.show_normalized else 1)
+            metrics[f'total_{self.feature_names[feat_idx]}'] = self.format_value(total, feat_idx)
+        
+        # Update plots
+        figures = self.create_plot(
+            self.latest_sample,
+            self.latest_observed_points,
+            self.latest_observed_mask,
+            self.latest_gradient_control_signal,
+            metrics
+        )
+        
+        return figures, metrics
+
+    def parse_data_points(self, df) -> Dict:
+        """Parse data points from DataFrame with columns: time,feature,value"""
+        data_dict = {}
+        if df is None or df.empty:
+            return data_dict
+
+        for _, row in df.iterrows():
+            # Skip if any required value is NaN
+            if pd.isna(row['time']) or pd.isna(row['feature']) or pd.isna(row['value']):
+                continue
+            try:
+                time_idx = int(row['time'])
+                feature_idx = int(row['feature'])
+                value = float(row['value'])
+                
+                if time_idx not in data_dict:
+                    data_dict[time_idx] = {}
+                data_dict[time_idx][feature_idx] = (value, 1.0)
+            except (ValueError, TypeError):
+                continue
+        return data_dict
+
+    def parse_point_groups(self, df) -> Dict:
+        """Parse point groups from DataFrame with columns: start,end,interval,feature,value,weight"""
+        data_dict = {}
+        if df is None or df.empty:
+            return data_dict
+            
+        for _, row in df.iterrows():
+            # Skip if any required value is NaN
+            if pd.isna(row['start']) or pd.isna(row['end']) or pd.isna(row['interval']) or \
+               pd.isna(row['feature']) or pd.isna(row['value']):
+                continue
+                
+            try:
+                start = int(row['start'])
+                end = int(row['end'])
+                interval = int(row['interval'])
+                feature = int(row['feature'])
+                value = float(row['value'])
+                weight = float(row.get('weight', 1.0)) if not pd.isna(row.get('weight')) else 1.0
+
+                for t in range(start, end + 1, interval):
+                    if 0 <= t < self.seq_length:
+                        if t not in data_dict:
+                            data_dict[t] = {}
+                        data_dict[t][feature] = (value, weight)
+            except (ValueError, TypeError):
+                continue
+        
+        return data_dict
+
+    def to_tensor(self, observed_points_dict, seq_length, feature_dim):
+        observed_points = torch.zeros((seq_length, feature_dim))
+        observed_weights = torch.zeros((seq_length, feature_dim))
+
+        for seq, feature_dict in observed_points_dict.items():
+            for feature, (value, weight) in feature_dict.items():
+                observed_points[seq, feature] = value
+                observed_weights[seq, feature] = weight
+
+        return observed_points, observed_weights
+
+    def apply_direct_edits(self, sample: np.ndarray, edit_params: Dict) -> np.ndarray:
+        """Apply direct edits to the sample array"""
+        edited_sample = sample.copy()
+        
+        if edit_params.get("enable_direct_area"):
+            areas = self.parse_area_selections(edit_params["direct_areas"])
+            for area in areas:
+                start, end, feat_idx, target = area
+                edited_sample[start:end, feat_idx] += target
+        edited_sample = np.clip(edited_sample, 0, 1)
+        return edited_sample
+
+    def parse_area_selections(self, area_text: str) -> List[Tuple]:
+        """Parse area selection text into (start, end, feature, target) tuples"""
+        areas = []
+        if not area_text.strip():
+            return areas
+
+        area_text = area_text.replace('\n', ';')
+
+        for line in area_text.strip().split(';'):
+            if not line.strip():
+                continue
+            try:
+                start, end, feat, target = map(float, line.strip().split(','))
+                areas.append((int(start), int(end), int(feat), target))
+            except (ValueError, IndexError):
+                continue
+        return areas
+
+    def apply_trending_mask(self, points: torch.Tensor, mask: torch.Tensor, consider_last_generated=False) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Apply trending functions as soft constraints through masks"""
+        if not self.trending_controls or self.latest_sample is None:
+            return points, mask
+
+        for start, end, feat_idx, func, confidence in self.trending_controls:
+            if start < 0 or end > self.seq_length or start >= end:
+                continue
+
+            # Generate x values normalized between 0 and 1 for the segment
+            x = np.linspace(0, 1, end - start)
+            
+            try:
+                # Calculate the function values
+                y = func(x)
+
+                # Scale the function output to 0-1 range
+                y = (y - np.min(y)) / (np.max(y) - np.min(y))
+                # points[start:end, feat_idx] = torch.tensor(y, dtype=points.dtype)
+                # mask[start:end, feat_idx] = max(mask[start:end, feat_idx], min(1.0, confidence * abs(
+                #     self.latest_sample[start:end, feat_idx] - y 
+                # )))  # Use lower weight for trending constraints
+
+            except Exception as e:
+                print(f"Error applying function: {e}")
+                continue
+
+        # Apply the trend as soft constraints
+        mask_zero = (mask[start:end, feat_idx] == 0)
+        points[start:end, feat_idx][mask_zero] = torch.tensor(y, dtype=points.dtype)[mask_zero]
+        mask[start:end, feat_idx][mask_zero] = torch.tensor(confidence * np.ones_like(y), dtype=mask.dtype)[mask_zero]
+
+        # mask[start:end, feat_idx][mask_zero] = torch.tensor((confidence * np.abs(self.latest_sample[start:end, feat_idx] - y)), dtype=mask.dtype)[mask_zero]
+        mask = mask.clamp(0, 1)
+
+        return points, mask
+    
+
+    def update_model(self, 
+                    figures: List[go.Figure],
+                    data_points: str,
+                    point_groups: str,
+                    enable_area_control: bool,
+                    area_selections: str,
+                    enable_auc: bool,
+                    auc_value: float,
+                    enable_peaks: bool,
+                    peak_points: str,
+                    peak_alpha: float,
+                    auc_weight: float,
+                    peak_weight: float,
+                    enable_trending: bool = False,
+                    enable_trending_with_diff: bool = False,
+                    trending_params: str = ""
+        ) -> Tuple[List[go.Figure], str, str, Dict]:
+
+        # Parse both point groups and individual data points
+        individual_points_dict = self.parse_data_points(data_points)
+        group_points_dict = self.parse_point_groups(point_groups)
+
+        # Merge dictionaries, giving preference to individual points
+        combined_points_dict = group_points_dict.copy()
+        for t, feat_dict in individual_points_dict.items():
+            if t not in combined_points_dict:
+                combined_points_dict[t] = {}
+            for f, v in feat_dict.items():
+                combined_points_dict[t][f] = v
+
+        # Convert to tensor
+        observed_points, observed_weights = self.to_tensor(
+            combined_points_dict, 
+            self.seq_length, 
+            self.feature_dim
+        )
+        observed_mask = observed_weights
+
+        # Parse peak points
+        peak_points_list = []
+        if enable_peaks and peak_points:
+            try:
+                peak_points_list = [int(x.strip()) for x in peak_points.split(',') if x.strip()]
+            except ValueError:
+                peak_points_list = []
+
+        # Apply trending control if enabled
+        if enable_trending and trending_params:
+            self.parse_trending_parameters(trending_params)
+            observed_points, observed_mask = self.apply_trending_mask(observed_points, observed_mask, consider_last_generated=enable_trending_with_diff)
+
+        # Build gradient control signal
+        # IMPORTANT
+        gradient_control_signal = {}
+        if enable_auc:
+            gradient_control_signal["auc"] = auc_value
+            gradient_control_signal["auc_weight"] = auc_weight
+        if enable_peaks:
+            gradient_control_signal.update({
+                "peak_points": peak_points_list,
+                "peak_alpha": peak_alpha,
+                "peak_weight": peak_weight
+            })
+
+        # Build model control signal
+        model_control_signal = {}
+        # if enable_area_control and area_selections:
+        #     areas = self.parse_area_selections(area_selections)
+        #     if areas:
+        #         model_control_signal["selected_areas"] = areas
+
+        # Run prediction
+        sample = self.trainer.predict_weighted_points(
+            observed_points, # (seq_length, feature_dim)
+            observed_mask, # (seq_length, feature_dim)
+            self.coef,  # fixed
+            self.stepsize, # fixed
+            self.sampling_steps, # fixed
+            # model_control_signal=model_control_signal,
+            gradient_control_signal=gradient_control_signal
+        )
+
+        # Store latest results
+        self.latest_sample = sample
+        self.latest_observed_points = observed_points
+        self.latest_observed_mask = observed_mask
+        self.latest_gradient_control_signal = gradient_control_signal
+        self.latest_model_control_signal = model_control_signal
+
+        # Calculate metrics
+        metrics = {
+            'show_normalized': self.show_normalized
+        }
+        for feat_idx in range(self.feature_dim):
+            total = np.sum(sample[:, feat_idx]) * (self.feature_scales[feat_idx] * 10 if not self.show_normalized else 1)
+            metrics[f'total_{self.feature_names[feat_idx]}'] = self.format_value(total, feat_idx)
+
+        # Update plots
+        figures = self.create_plot(sample, observed_points, observed_mask, gradient_control_signal, metrics)
+
+        return figures, data_points, point_groups, metrics
+
+
+    def update_additional_edit(
+        self,
+                    enable_direct_area: bool,
+                    direct_areas: str):
+        # Apply direct edits if enabled
+        if enable_direct_area:
+            sample = self.apply_direct_edits(self.latest_sample, {
+                "enable_direct_area": enable_direct_area,
+                "direct_areas": direct_areas
+            })
+        else:
+            sample = self.latest_sample
+
+        # Calculate metrics
+        metrics = {
+            'show_normalized': self.show_normalized
+        }
+        for feat_idx in range(self.feature_dim):
+            total = np.sum(sample[:, feat_idx]) * (self.feature_scales[feat_idx] * 10 if not self.show_normalized else 1)
+            metrics[f'total_{self.feature_names[feat_idx]}'] = self.format_value(total, feat_idx)
+
+        # Update plots
+        figures = self.create_plot(
+            sample,
+            self.latest_observed_points,
+            self.latest_observed_mask,
+            self.latest_gradient_control_signal,
+            metrics
+        )
+
+        return figures, metrics
+
+
+    def apply_frequency_filter(self, signal: np.ndarray) -> np.ndarray:
+        """Apply FFT-based frequency filtering using the current band multipliers"""
+        # Get FFT of the signal
+        fft = np.fft.fft(signal)
+        freqs = np.fft.fftfreq(len(signal))
+        
+        # Split frequencies into 5 bands
+        # Exclude DC component (0 frequency) from bands
+        pos_freqs = freqs[1:len(freqs)//2]
+        freq_ranges = np.array_split(pos_freqs, 5)
+        
+        # Apply band multipliers
+        filtered_fft = fft.copy()
+        
+        # Handle DC component separately (lowest frequency)
+        filtered_fft[0] *= self.freq_bands[4]  # Apply very low freq multiplier to DC
+        
+        # Apply multipliers to each frequency band
+        for i, freq_range in enumerate(freq_ranges):
+            # Get indices for this frequency band
+            band_mask = np.logical_and(
+                freqs >= freq_range[0],
+                freqs <= freq_range[-1]
+            )
+            
+            # Apply multiplier to positive and negative frequencies symmetrically
+            filtered_fft[band_mask] *= self.freq_bands[4-i]
+            filtered_fft[np.flip(band_mask)] *= self.freq_bands[4-i]
+        
+        # Convert back to time domain
+        filtered_signal = np.real(np.fft.ifft(filtered_fft))
+        
+        return filtered_signal
+    
+    
+    def update_frequency_bands(self, band_idx: int, value: float) -> Tuple[List[go.Figure], Dict]:
+        """Update a frequency band multiplier and recompute the filtered signal"""
+        if self.latest_sample is None:
+            return [], {}
+        
+        # Update the specified band multiplier
+        self.freq_bands[band_idx] = value
+        
+        # Apply frequency filtering to each feature
+        filtered_sample = self.latest_sample.copy()
+        for feat_idx in range(self.feature_dim):
+            filtered_sample[:, feat_idx] = self.apply_frequency_filter(
+                self.latest_sample[:, feat_idx]
+            )
+        
+        # Ensure values remain in valid range
+        filtered_sample = np.clip(filtered_sample, 0, 1)
+        
+        # Calculate metrics
+        metrics = {
+            'show_normalized': self.show_normalized,
+            'frequency_bands': self.freq_bands.tolist()
+        }
+        for feat_idx in range(self.feature_dim):
+            total = np.sum(filtered_sample[:, feat_idx]) * (self.feature_scales[feat_idx] * 10 if not self.show_normalized else 1)
+            metrics[f'total_{self.feature_names[feat_idx]}'] = self.format_value(total, feat_idx)
+        
+        # Update plots
+        figures = self.create_plot(
+            filtered_sample,
+            self.latest_observed_points,
+            self.latest_observed_mask,
+            self.latest_gradient_control_signal,
+            metrics
+        )
+        
+        return figures, metrics
+
+    def parse_trending_parameters(self, trending_text: str) -> List[Tuple]:
+        """Parse trending control parameters into (start, end, feature, function) tuples"""
+        trending_params = []
+        if not trending_text.strip():
+            return trending_params
+
+        trending_text = trending_text.replace('\n', ';')
+
+        for line in trending_text.strip().split(';'):
+            if not line.strip():
+                continue
+            try:
+                # Split by comma and handle the function part separately
+                parts = line.strip().split(',', 4)
+                if len(parts) != 5:
+                    continue
+
+                start, end, feat = map(int, parts[:3])
+                function_str = parts[3].strip()
+                confidence = float(parts[4])
+                # Convert the function string to a callable
+                try:
+                    func = self.function_parser.string_to_function(function_str)
+                    trending_params.append((start, end, feat, func, confidence))
+                except ValueError as e:
+                    print(f"Error parsing function '{function_str}': {e}")
+                    continue
+
+            except (ValueError, IndexError):
+                continue
+        self.trending_controls = trending_params  # Store the parsed parameters
+        return trending_params
+
+
+def create_gradio_interface(editor: TimeSeriesEditor):
+    with gr.Blocks() as app:
+        gr.Markdown("# Time Series Editor")
+        gr.Markdown("## Instruction: Scroll Down + Click `Update Figure` [~20s]")
+
+        metrics_display = gr.JSON(label="Metrics", value={})
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                # with Tab():
+                # Scaling Parameters Section
+                # with gr.Group():
+
+                gr.Markdown("## Scaling Parameters")
+                with gr.Accordion("Open for More Detail", open=False):
+                    revenue_scale = gr.Number(
+                        label="Revenue Scale ($ per 0.1 in model)",
+                        value=1000000
+                    )
+                    download_scale = gr.Number(
+                        label="Download Scale (downloads per 0.1 in model)",
+                        value=100000
+                    )
+                    au_scale = gr.Number(
+                        label="Active Users Scale (users per 0.1 in model)",
+                        value=10000
+                    )
+                    show_normalized = gr.Checkbox(
+                        label="Show Normalized Values (0-1 scale)",
+                        value=True
+                    )
+                    update_scaling_btn = gr.Button("Update Scaling")
+
+                # TS Section
+                gr.Markdown("## Time Series Control Panel")
+                with gr.Accordion("Open for More Detail"):
+                    with gr.Group():
+                        gr.Markdown("### Anchor Point Control")
+                        data_points_df = gr.Dataframe(
+                            headers=["time", "feature", "value"],
+                            datatype=["number", "number", "number"],
+                            # label="Anchor Point Control",
+                            value=[[0, 0, 0.04], [2, 0, 0.58], [6, 0, 0.27], [58, 0, 1.0], [60, 0, 0.5]],
+                            col_count=(3, "fixed"),  # Fix number of columns
+                            interactive=True
+                        )
+                        add_data_point_btn = gr.Button("Add Data Point")
+
+                        def add_data_point(df):
+                            new_row = pd.DataFrame([[None, 0, None]], 
+                                                 columns=["time", "feature", "value"])
+                            return pd.concat([df, new_row], ignore_index=True)
+
+                        add_data_point_btn.click(
+                            fn=add_data_point,
+                            inputs=[data_points_df],
+                            outputs=[data_points_df]
+                        )
+
+                    with gr.Group():
+                        gr.Markdown("### Group of Anchor Point Control")
+                        point_groups_df = gr.Dataframe(
+                            headers=["start", "end", "interval", "feature", "value", "weight"],
+                            datatype=["number", "number", "number", "number", "number", "number"],
+                            # label="Group of Anchor Point Control",
+                            value=[[0, 50, 10, 0, 0.5, 0.1], [100, 150, 50, 0, 0.1, 0.5]],
+                            col_count=(6, "fixed"),  # Fix number of columns
+                            interactive=True
+                        )
+                        add_point_group_btn = gr.Button("Add Point Group")
+
+                        def add_point_group(df):
+                            new_row = pd.DataFrame([[None, None, None, 0, None, None]], 
+                                                 columns=["start", "end", "interval", "feature", "value", "weight"])
+                            return pd.concat([df, new_row], ignore_index=True)
+
+                        add_point_group_btn.click(
+                            fn=add_point_group,
+                            inputs=[point_groups_df],
+                            outputs=[point_groups_df]
+                        )
+
+                    with gr.Group():
+                    # with gr.Tab("Trending Control"):
+                        gr.Markdown("### Trending Control")
+                        gr.Markdown("""
+                        Enter trending control parameters in the format:
+                        ```
+                        start_time,end_time,feature,function,confidence
+                        ```
+                        Examples:
+                        - Linear trend: `0,100,0,x`
+                        - Sine wave: `0,100,0,sin(2*pi*x)`
+                        - Exponential: `0,100,0,exp(-x)`
+                        
+                        Separate multiple trends with semicolons.
+                        """)
+                        enable_trending_control = gr.Checkbox(label="Enable Trending Control", value=True)
+                        enable_trending_control_with_diff = gr.Checkbox(label="Consider Last Generated", value=False)
+                        trending_control = gr.Textbox(
+                            label="Trending Control Parameters",
+                            lines=2,
+                            placeholder="Enter parameters: start_time,end_time,feature,function,condifdence; separated by semicolons",
+                            value="200,250,0,sin(2*pi*x),0.2"
+                        )
+
+                    # Area Control Parameters
+                    with gr.Group(visible=False):
+                        gr.Markdown("### Area Control")
+                        enable_area_control = gr.Checkbox(label="Enable Area Control", value=False)
+                        area_selections = gr.Textbox(
+                            label="Area Selections (format: start_time,end_time,feature,target_value)",
+                            lines=2,
+                            placeholder="Enter areas: start,end,feature,target; separated by semicolons",
+                            
+                        )
+                        
+                    # AUC Parameters
+                    gr.Markdown("### Statistics Control")
+                    enable_auc = gr.Checkbox(label="Enable Total Sum Control", value=True)
+                    auc_input = gr.Number(label="Target Sum Value", value=-150)
+                    auc_weight_input = gr.Number(label="Sum Weight", value=10.0)
+                    
+                    # Peak Parameters
+                    with gr.Group(visible=False):
+                        gr.Markdown("### Peak Control")
+                        enable_peaks = gr.Checkbox(label="Enable Peak Control", value=False)
+                        peak_points_input = gr.Textbox(label="Peak Points (comma-separated)", value="100,200")
+                        peak_alpha_input = gr.Number(label="Peak Alpha", value=10)
+                        peak_weight_input = gr.Number(label="Peak Weight", value=1.0)
+                        
+                    update_model_btn = gr.Button("Update Figure")
+
+                gr.Markdown("## Extend Edit", visible=False)
+                with gr.Tab("Range Shift", visible=False):
+                    # gr.Markdown("### Direct Edit Control")
+                    enable_direct_area = gr.Checkbox(label="Enable Direct Edits", value=False) # range shift
+                    direct_areas = gr.Textbox(
+                        label="Direct Edit Areas (format: start_time,end_time,feature,delta)",
+                        lines=2,
+                        placeholder="Enter areas: start,end,feature,delta; separated by semicolons",
+                        value="150,200,0,-0.1"
+                    )
+                    
+                    update_additional_btn = gr.Button("Update Additional Edit")
+
+                # with gr.Tab("Trending Control"):
+                #     gr.Markdown("### Trending Control")
+                #     gr.Markdown("""
+                #     Enter trending control parameters in the format:
+                #     ```
+                #     start_time,end_time,feature,function
+                #     ```
+                #     Examples:
+                #     - Linear trend: `0,100,0,x`
+                #     - Sine wave: `0,100,0,sin(2*pi*x)`
+                #     - Exponential: `0,100,0,exp(-x)`
+                    
+                #     Separate multiple trends with semicolons.
+                #     """)
+                #     enable_trending_control = gr.Checkbox(label="Enable Trending Control", value=False)
+                #     enable_trending_control_with_diff = gr.Checkbox(label="Consider Last Generated", value=False)
+                #     trending_control = gr.Textbox(
+                #         label="Trending Control Parameters",
+                #         lines=2,
+                #         placeholder="Enter parameters: start_time,end_time,feature,function,condifdence; separated by semicolons",
+                #         value="0,100,0,sin(2*pi*x),0.3"
+                #     )
+
+                # with gr.Tab("Frequency Controls", visible=False):
+                with gr.Group(visible=False):
+                    gr.Markdown("Adjust multipliers for different frequency bands (0-2)")
+                    freq_bands = [
+                        gr.Slider(
+                            minimum=0, maximum=2, step=0.1, value=1.0,
+                            label=f"Band {i+1}: {'Very High' if i==0 else 'High' if i==1 else 'Mid' if i==2 else 'Low' if i==3 else 'Very Low'} Freq",
+                        ) for i in range(5)
+                    ]
+
+                gr.Markdown("### Feature Index Reference:")
+                for idx, name in enumerate(editor.feature_names):
+                    gr.Markdown(f"- {idx}: {name}")
+
+            with gr.Column(scale=1.2):
+                gr.Markdown("""
+                ### Plot Legend
+                - **Points with Error Bars**: Observed values where:
+                    - Point position = observed value
+                    - Error bar size = uncertainty (inversely proportional to weight)
+                - **Green Line**: Model prediction
+                - **Vertical Red Lines**: Peak points (if enabled)
+                """)
+                plots = [gr.Plot() for _ in range(editor.feature_dim)]
+                # - **Shaded Area**: General prediction uncertainty
+
+        def update_scaling_callback(revenue_scale, download_scale, au_scale, show_normalized):
+            figs, metrics = editor.update_scaling(
+                revenue_scale,
+                download_scale,
+                au_scale,
+                show_normalized
+            )
+            return [*figs, metrics]
+
+        def update_model_callback(
+            data_points_df, 
+            point_groups_df,
+            enable_area_control, 
+            area_selections,
+            enable_auc, 
+            auc, 
+            auc_weight,
+            enable_peaks, 
+            peak_points, 
+            peak_alpha, 
+            peak_weight,
+            enable_trending,
+            enable_trending_with_diff,
+            trending_params
+        ):
+            figs, _, _, metrics = editor.update_model(
+                plots,
+                data_points_df,
+                point_groups_df,
+                enable_area_control,
+                area_selections,
+                enable_auc,
+                auc,
+                enable_peaks,
+                peak_points,
+                peak_alpha,
+                auc_weight,
+                peak_weight,
+                enable_trending,
+                enable_trending_with_diff,
+                trending_params
+            )
+            return [*figs, metrics]
+
+        # Update the click handler
+        update_model_btn.click(
+            fn=update_model_callback,
+            inputs=[
+                data_points_df,
+                point_groups_df,
+                enable_area_control, 
+                area_selections,
+                enable_auc, 
+                auc_input, 
+                auc_weight_input,
+                enable_peaks, 
+                peak_points_input, 
+                peak_alpha_input, 
+                peak_weight_input,
+                enable_trending_control,
+                enable_trending_control_with_diff,
+                trending_control
+            ],
+            outputs=[*plots, metrics_display]
+        )
+
+        
+        def update_additional_callback(enable_direct_area, direct_areas):
+            figs, metrics = editor.update_additional_edit(
+                enable_direct_area,
+                direct_areas
+            )
+            return [*figs, metrics]
+
+        def update_freq_band(band_idx, value):
+            figs, metrics = editor.update_frequency_bands(band_idx, value)
+            return [*figs, metrics]
+        
+        update_scaling_btn.click(
+            fn=update_scaling_callback,
+            inputs=[
+                revenue_scale,
+                download_scale,
+                au_scale,
+                show_normalized
+            ],
+            outputs=[*plots, metrics_display]
+        )
+
+        update_additional_btn.click(
+            fn=update_additional_callback,
+            inputs=[enable_direct_area, direct_areas],
+            outputs=[*plots, metrics_display]
+        )
+
+        # Add event handlers for frequency band sliders
+        for i, slider in enumerate(freq_bands):
+            slider.change(
+                fn=update_freq_band,
+                inputs=[gr.Number(value=i, visible=False), slider],
+                outputs=[*plots, metrics_display]
+            )
+
+        # app.load(
+        #     fn=update_model_callback,
+        #     inputs=[
+        #         data_points_df,
+        #         point_groups_df,
+        #         enable_area_control, 
+        #         area_selections,
+        #         enable_auc, 
+        #         auc_input, 
+        #         auc_weight_input,
+        #         enable_peaks, 
+        #         peak_points_input, 
+        #         peak_alpha_input, 
+        #         peak_weight_input,
+        #         enable_trending_control,
+        #         enable_trending_control_with_diff,
+        #         trending_control
+        #     ],
+        #     outputs=[*plots, metrics_display]
+        # )
+
+    return app
+
+
+class FunctionParser:
+    def __init__(self):
+        # Define available mathematical functions and constants
+        self.math_functions = {
+            'sin': np.sin,
+            'cos': np.cos,
+            'tan': np.tan,
+            'exp': np.exp,
+            'log': np.log,
+            'sqrt': np.sqrt,
+            'abs': np.abs,
+            'pow': np.power,
+            'pi': np.pi,
+            'e': np.e,
+            'asin': np.arcsin,
+            'acos': np.arccos,
+            'atan': np.arctan,
+            'sinh': np.sinh,
+            'cosh': np.cosh,
+            'tanh': np.tanh
+        }
+
+    def validate_expression(self, expression: str) -> bool:
+        """
+        Validate the mathematical expression for basic syntax errors.
+        """
+        # Check for balanced parentheses
+        if expression.count('(') != expression.count(')'):
+            raise ValueError("Unbalanced parentheses in expression")
+        
+        # Check for invalid characters
+        valid_chars = set('0123456789.+-*/()^ xXepi,')
+        valid_chars.update(''.join(self.math_functions.keys()))
+        if not all(c in valid_chars or c.isspace() for c in expression.lower()):
+            raise ValueError("Expression contains invalid characters")
+        
+        return True
+
+    def preprocess_expression(self, expression: str) -> str:
+        """
+        Preprocess the expression to handle various input formats.
+        """
+        # Remove whitespace
+        expression = expression.replace(' ', '')
+        
+        # Convert ^ to ** for exponentiation
+        expression = expression.replace('^', '**')
+        
+        # Ensure multiplication is explicit
+        expression = re.sub(r'(\d+)([a-zA-Z])', r'\1*\2', expression)
+        expression = re.sub(r'(\))([\w])', r'\1*\2', expression)
+        
+        # Replace X with x for consistency
+        expression = expression.lower()
+        
+        return expression
+
+    def string_to_function(self, expression: str) -> Callable[[Union[float, np.ndarray]], Union[float, np.ndarray]]:
+        """
+        Convert a string mathematical expression to a callable function.
+        
+        Args:
+            expression (str): Mathematical expression (e.g., "sin(x) + x^2")
+            
+        Returns:
+            Callable: A function that takes x as input and returns the evaluated result
+            
+        Example:
+            >>> f = string_to_function("sin(x) + x^2")
+            >>> f(0.5)
+            0.729321...
+        """
+        # Validate and preprocess the expression
+        self.validate_expression(expression)
+        processed_expr = self.preprocess_expression(expression)
+        
+        # Create the function namespace
+        namespace = self.math_functions.copy()
+        
+        try:
+            # Create the lambda function
+            func = eval(f"lambda x: {processed_expr}", namespace)
+            
+            # Test the function with a simple input
+            test_value = 1.0
+            try:
+                func(test_value)
+            except Exception as e:
+                raise ValueError(f"Invalid function: {str(e)}")
+            
+            return func
+            
+        except SyntaxError as e:
+            raise ValueError(f"Invalid expression syntax: {str(e)}")
+        except Exception as e:
+            raise ValueError(f"Error creating function: {str(e)}")
+
+    @staticmethod
+    def demonstrate_usage():
+        """
+        Demonstrate various uses of the function parser.
+        """
+        parser = FunctionParser()
+        
+        # Test cases
+        test_expressions = [
+            "x^2 + 2*x + 1",
+            "sin(x) + cos(x)",
+            "exp(-x^2)",
+            "log(x + 1)",
+            "sqrt(1 - x^2)",
+        ]
+        
+        print("Testing various mathematical expressions:")
+        x_test = 0.5
+        
+        for expr in test_expressions:
+            try:
+                print(f"\nExpression: {expr}")
+                func = parser.string_to_function(expr)
+                result = func(x_test)
+                print(f"f({x_test}) = {result}")
+                
+                # Test with numpy array
+                x_array = np.linspace(0, 1, 5)
+                results = func(x_array)
+                print(f"f(array) = {results}")
+                
+            except Exception as e:
+                print(f"Error: {str(e)}")
+
+# Example usage:
+if __name__ == "__main__":
+    # Initialize with example data points
+    # example_data_points = "0,0,0.04;2,0,0.58;6,0,0.27;58,0,1.0;-1,0,0.05"
+    
+    import os
+    import torch
+    import numpy as np
+    from engine.solver import Trainer
+    from utils.io_utils import load_yaml_config, instantiate_from_config
+
+    # assert torch.cuda.is_available(), "CUDA must be available"
+    class Parameters:
+        def __init__(self) -> None:
+            self.gpu = 0
+            self.config_path = "./config/modified/revenue-baseline-365.yaml"
+            # self.config_path = "config/modified/96/fmri.yaml"
+            # self.config_path = "./config/control/revenue-baseline-sine.yaml"
+            # self.save_dir = (
+            #     "../../../data/" + os.path.basename(self.config_path).split(".")[0]
+            # )
+            self.mode = "infill"
+            self.missing_ratio = 0.95
+            self.milestone = "10"
+            # os.makedirs(self.save_dir, exist_ok=True)
+
+    os.environ["WANDB_ENABLED"] = "false"
+    # os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+    # print working directory
+    print(os.getcwd())
+    args = Parameters()
+    configs = load_yaml_config(args.config_path)
+    # device = torch.device('cpu')
+    device = torch.device(f"cuda:0" if torch.cuda.is_available() else "cpu")
+
+    # dl_info = build_dataloader_cond(configs, args)
+    model = instantiate_from_config(configs["model"]).to(device)
+    trainer = Trainer(config=configs, args=args, model=model, dataloader={
+        "dataloader": []
+    })
+
+    trainer.load(args.milestone, from_folder="./weight") #, from_folder="../../../data/ckpt_baseline_sine_240"), from_folder="./data/weight_365"
+    # dataloader, dataset = dl_info["dataloader"], dl_info["dataset"]
+    coef = configs["dataloader"]["test_dataset"]["coefficient"]
+    stepsize = configs["dataloader"]["test_dataset"]["step_size"]
+    sampling_steps = configs["dataloader"]["test_dataset"]["sampling_steps"]
+    seq_length = configs["dataloader"]["test_dataset"]["params"]["window"]
+    feature_dim = 3
+    print(f"seq_length: {seq_length}, feature_dim: {feature_dim}")
+
+    # Initialize your trainer, configs, and dataset here
+    editor = TimeSeriesEditor(seq_length, feature_dim, trainer)
+    editor.coef = coef
+    editor.stepsize = stepsize
+    editor.sampling_steps = sampling_steps
+    
+    app = create_gradio_interface(editor)
+    # app.launch(server_name="0.0.0.0", server_port=8888, share=True)
+    app.launch(show_api=False)

config/backup/revenue-1.yaml

@@ -0,0 +1,76 @@
+model:
+  target: Models.interpretable_diffusion.gaussian_diffusion.Diffusion_TS
+  params:
+    seq_length: 240
+    feature_size: 3
+    n_layer_enc: 8
+    n_layer_dec: 5
+    d_model: 128  # 4 X 16
+    timesteps: 1000 # diffusion timesteps
+    sampling_timesteps: 200
+    loss_type: 'l2'
+    beta_schedule: 'cosine'
+    n_heads: 8
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 1385
+  results_folder: ../../../data/Checkpoints_revenue-1
+  gradient_accumulate_every: 1
+  save_cycle: 277  # max_epochs // 5
+  ema:
+    decay: 0.995
+    update_interval: 10
+
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 200
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 300 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # Set to rate < 1 if training conditional generation
+      # data_root: ./Data/datasets/stock_data.csv
+      data_root: ../../../data/daily.csv
+      window: 240  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # rate
+      data_root: ../../../data/daily.csv
+      window: 240  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: test
+      style: separate
+      distribution: geometric
+
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 200
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/backup/revenue-2.yaml

@@ -0,0 +1,77 @@
+model:
+  target: Models.interpretable_diffusion.gaussian_diffusion.Diffusion_TS
+  params:
+    seq_length: 240
+    feature_size: 3
+    n_layer_enc: 6
+    n_layer_dec: 4
+    d_model: 128  # 4 X 16
+    timesteps: 500 # diffusion timesteps
+    sampling_timesteps: 500
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 8
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 2770
+  results_folder: ../../../data/Checkpoints_revenue-2
+  gradient_accumulate_every: 1
+  save_cycle: 277  # max_epochs // 5
+  ema:
+    decay: 0.995
+    update_interval: 10
+
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 200
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 300 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # Set to rate < 1 if training conditional generation
+      # data_root: ./Data/datasets/stock_data.csv
+      data_root: ../../../data/daily.csv
+      window: 240  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # rate
+      data_root: ../../../data/daily.csv
+      window: 240  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: test
+      style: separate
+      # distribution: geometric
+      distribution: uniform
+
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 200
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/backup/revenue-3xl.yaml

@@ -0,0 +1,77 @@
+model:
+  target: Models.interpretable_diffusion.gaussian_diffusion.Diffusion_TS
+  params:
+    seq_length: 480
+    feature_size: 3
+    n_layer_enc: 8
+    n_layer_dec: 4
+    d_model: 256  # 4 X 16
+    timesteps: 2000 # diffusion timesteps
+    sampling_timesteps: 500
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 16
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 5540
+  results_folder: ../../../data/Checkpoints_revenue-3xl
+  gradient_accumulate_every: 1
+  save_cycle: 554  # max_epochs // 5
+  ema:
+    decay: 0.995
+    update_interval: 10
+
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 200
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 300
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # Set to rate < 1 if training conditional generation
+      # data_root: ./Data/datasets/stock_data.csv
+      data_root: ../../../data/daily.csv
+      window: 480  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # rate
+      data_root: ../../../data/daily.csv
+      window: 480  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: test
+      style: separate
+      # distribution: geometric
+      distribution: uniform
+
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 200
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/backup/revenue-baseline.yaml

@@ -0,0 +1,82 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 240
+    feature_size: 3
+    n_layer_enc: 6
+    n_layer_dec: 4
+    d_model: 128  # 4 X 16
+    timesteps: 500 # diffusion timesteps
+    sampling_timesteps: 500
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 8
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+      # - classifier-based-sum-control
+      # - classifier-free-sum-control
+      # - range-wise-peak-control
+
+solver:
+  base_lr: 2.0e-5
+  max_epochs: 2230 # 11150
+  results_folder: ../../../data/ckpt_baseline
+  gradient_accumulate_every: 2
+  save_cycle: 223 # 1115  # max_epochs // 5
+  ema:
+    decay: 0.995
+    update_interval: 10
+
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.65
+      patience: 200
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 100
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # Set to rate < 1 if training conditional generation
+      # data_root: ./Data/datasets/stock_data.csv
+      data_root: ../../../data/daily.csv
+      window: 240  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # rate
+      data_root: ../../../data/daily.csv
+      window: 240  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: test
+      style: separate
+      # distribution: geometric
+      distribution: uniform
+      missing_ratio: 0.5
+
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 100
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/backup/revenue-test.yaml

@@ -0,0 +1,82 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 240
+    feature_size: 3
+    n_layer_enc: 6
+    n_layer_dec: 4
+    d_model: 128  # 4 X 16
+    timesteps: 500 # diffusion timesteps
+    sampling_timesteps: 500
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 8
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+      # - classifier-based-sum-control
+      # - classifier-free-sum-control
+      # - range-wise-peak-control
+
+solver:
+  base_lr: 2.0e-5
+  max_epochs: 2230 # 11150
+  results_folder: ../../../data/ckpt_baseline
+  gradient_accumulate_every: 2
+  save_cycle: 223 # 1115  # max_epochs // 5
+  ema:
+    decay: 0.995
+    update_interval: 10
+
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.65
+      patience: 200
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 100
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # Set to rate < 1 if training conditional generation
+      # data_root: ./Data/datasets/stock_data.csv
+      data_root: ../../../data/daily.csv
+      window: 240  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # rate
+      data_root: ../../../data/daily.csv
+      window: 240  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: test
+      style: separate
+      # distribution: geometric
+      distribution: uniform
+      missing_ratio: 0.5
+
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 500
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/backup/revenue.yaml

@@ -0,0 +1,76 @@
+model:
+  target: Models.interpretable_diffusion.gaussian_diffusion.Diffusion_TS
+  params:
+    seq_length: 240
+    feature_size: 3
+    n_layer_enc: 5
+    n_layer_dec: 2
+    d_model: 64  # 4 X 16
+    timesteps: 500 # diffusion timesteps
+    sampling_timesteps: 500
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 2770
+  results_folder: ../../../data/Checkpoints_revenue
+  gradient_accumulate_every: 2
+  save_cycle: 554  # max_epochs // 5
+  ema:
+    decay: 0.995
+    update_interval: 10
+
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 2000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 300 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.7  # Set to rate < 1 if training conditional generation
+      # data_root: ./Data/datasets/stock_data.csv
+      data_root: ../../../data/daily.csv
+      window: 240  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.9  # rate
+      data_root: ../../../data/daily.csv
+      window: 240  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: test
+      style: separate
+      distribution: geometric
+
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 200
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/config.yaml

@@ -0,0 +1,40 @@
+model:
+  target: Models.interpretable_diffusion.gaussian_diffusion.Diffusion_TS
+  params:
+    seq_length: 160
+    feature_size: 5
+    n_layer_enc: 1
+    n_layer_dec: 2
+    d_model: 64  # 4 X 16
+    timesteps: 200
+    sampling_timesteps: 200
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.1
+    resid_pd: 0.1
+    kernel_size: 5
+    padding_size: 2
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 1000
+  results_folder: ./Checkpoints_syn
+  gradient_accumulate_every: 2
+  save_cycle: 100  # max_epochs // 10
+  ema:
+    decay: 0.99
+    update_interval: 10
+
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 200
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 100 
+      verbose: False

config/control/revenue-baseline-180.yaml

@@ -0,0 +1,82 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 180
+    feature_size: 3
+    n_layer_enc: 6
+    n_layer_dec: 4
+    d_model: 128  # 4 X 16
+    timesteps: 500 # diffusion timesteps
+    sampling_timesteps: 500
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 8
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+      # - classifier-based-sum-control
+      # - classifier-free-sum-control
+      # - range-wise-peak-control
+
+solver:
+  base_lr: 2.0e-5
+  max_epochs: 2230 # 11150
+  results_folder: ../../../data/ckpt_baseline
+  gradient_accumulate_every: 2
+  save_cycle: 223 # 1115  # max_epochs // 5
+  ema:
+    decay: 0.995
+    update_interval: 10
+
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.65
+      patience: 200
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 100
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # Set to rate < 1 if training conditional generation
+      # data_root: ./Data/datasets/stock_data.csv
+      data_root: ../../../data/daily.csv
+      window: 180  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # rate
+      data_root: ../../../data/daily.csv
+      window: 180  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: test
+      style: separate
+      # distribution: geometric
+      distribution: uniform
+      missing_ratio: 0.5
+
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 100
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/control/revenue-baseline-365-ma.yaml

@@ -0,0 +1,83 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 365
+    feature_size: 3
+    n_layer_enc: 6
+    n_layer_dec: 4
+    d_model: 128  # 4 X 16
+    timesteps: 500 # diffusion timesteps
+    sampling_timesteps: 200
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 8
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+    moving_average: True
+      # - classifier-based-sum-control
+      # - classifier-free-sum-control
+      # - range-wise-peak-control
+
+solver:
+  base_lr: 2.0e-5
+  max_epochs: 2230 # 11150
+  results_folder: ../../../data/ckpt_ma
+  gradient_accumulate_every: 2
+  save_cycle: 223 # 1115  # max_epochs // 5
+  ema:
+    decay: 0.995
+    update_interval: 10
+
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.65
+      patience: 200
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 100
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # Set to rate < 1 if training conditional generation
+      # data_root: ./Data/datasets/stock_data.csv
+      data_root: ../../../data/daily.csv
+      window: 365  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # rate
+      data_root: ../../../data/daily.csv
+      window: 365  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: test
+      style: separate
+      # distribution: geometric
+      distribution: uniform
+      missing_ratio: 0.5
+
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 100
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/control/revenue-baseline-365.yaml

@@ -0,0 +1,82 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 365
+    feature_size: 3
+    n_layer_enc: 6
+    n_layer_dec: 4
+    d_model: 128  # 4 X 16
+    timesteps: 500 # diffusion timesteps
+    sampling_timesteps: 200
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 8
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+      # - classifier-based-sum-control
+      # - classifier-free-sum-control
+      # - range-wise-peak-control
+
+solver:
+  base_lr: 2.0e-5
+  max_epochs: 2230 # 11150
+  results_folder: ../../../data/ckpt_baseline
+  gradient_accumulate_every: 2
+  save_cycle: 223 # 1115  # max_epochs // 5
+  ema:
+    decay: 0.995
+    update_interval: 10
+
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.65
+      patience: 200
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 100
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # Set to rate < 1 if training conditional generation
+      # data_root: ./Data/datasets/stock_data.csv
+      data_root: ../../../data/daily.csv
+      window: 365  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # rate
+      data_root: ../../../data/daily.csv
+      window: 365  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: test
+      style: separate
+      # distribution: geometric
+      distribution: uniform
+      missing_ratio: 0.5
+
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 100
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/control/revenue-baseline-sine.yaml

@@ -0,0 +1,82 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 240
+    feature_size: 3
+    n_layer_enc: 6
+    n_layer_dec: 4
+    d_model: 128  # 4 X 16
+    timesteps: 500 # diffusion timesteps
+    sampling_timesteps: 500
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 8
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+      # - classifier-based-sum-control
+      # - classifier-free-sum-control
+      # - range-wise-peak-control
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 223 # 11150
+  results_folder: ../../../data/ckpt_baseline_sine
+  gradient_accumulate_every: 2
+  save_cycle: 223 # 1115  # max_epochs // 5
+  ema:
+    decay: 0.995
+    update_interval: 10
+
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.65
+      patience: 20
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 10
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.sine_dataset.SineDataset
+    params:
+      name: revenue
+      proportion: 0.8  # Set to rate < 1 if training conditional generation
+      # data_root: ./Data/datasets/stock_data.csv
+      data_root: ../../../data/daily.csv
+      window: 240  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # rate
+      data_root: ../../../data/daily.csv
+      window: 240  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: test
+      style: separate
+      # distribution: geometric
+      distribution: uniform
+      missing_ratio: 0.5
+
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 100
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/control/revenue-extend.yaml

@@ -0,0 +1,83 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 240
+    feature_size: 3
+    n_layer_enc: 6
+    n_layer_dec: 4
+    d_model: 128  # 4 X 16
+    timesteps: 500 # diffusion timesteps
+    sampling_timesteps: 500
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 8
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+      # - classifier-based-sum-control
+      # - classifier-free-sum-control
+      # - range-wise-peak-control
+      # - regression-based-sum-control
+
+solver:
+  base_lr: 2.0e-5
+  max_epochs: 11150
+  results_folder: ../../../data/ckpt_baseline_extend
+  gradient_accumulate_every: 2
+  save_cycle: 1115  # max_epochs // 5
+  ema:
+    decay: 0.995
+    update_interval: 10
+
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.65
+      patience: 200
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 100
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.ControlRevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # Set to rate < 1 if training conditional generation
+      # data_root: ./Data/datasets/stock_data.csv
+      data_root: ../../../data/daily.csv
+      window: 240  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # rate
+      data_root: ../../../data/daily.csv
+      window: 240  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: test
+      style: separate
+      # distribution: geometric
+      distribution: uniform
+      missing_ratio: 0.5
+
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 100
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/csdi/energy.yaml

@@ -0,0 +1,75 @@
+model:
+  target: models.CSDI.tiffusion.Tiffusion
+  params:
+    seq_length: 24
+    feature_size: 28
+    n_layer_enc: 4
+    n_layer_dec: 3
+    d_model: 96  # 4 X 24
+    timesteps: 1000
+    sampling_timesteps: 1000
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+
+solver:
+  base_lr: 1.0e-3
+  max_epochs: 25000
+  results_folder: ../../../data/CSDI/Checkpoints_energy
+  gradient_accumulate_every: 2
+  save_cycle: 2500  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 5000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.CustomDataset
+    params:
+      name: energy
+      proportion: 1.0  # Set to rate < 1 if training conditional generation
+      data_root: ./data/energy_data.csv
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.CustomDataset
+    params:
+      name: energy
+      proportion: 0.9  # rate
+      data_root: ./data/energy_data.csv
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: test
+      style: separate
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 250
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/csdi/fmri.yaml

@@ -0,0 +1,74 @@
+model:
+  target: models.CSDI.tiffusion.Tiffusion
+  params:
+    seq_length: 24
+    feature_size: 50
+    n_layer_enc: 4
+    n_layer_dec: 4
+    d_model: 96  # 4 X 24
+    timesteps: 1000
+    sampling_timesteps: 1000
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 5
+    padding_size: 2
+
+solver:
+  base_lr: 1.0e-3
+  max_epochs: 15000
+  results_folder: ../../../data/CSDI/Checkpoints_fmri
+  gradient_accumulate_every: 2
+  save_cycle: 1500  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 3000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.fMRIDataset
+    params:
+      name: fMRI
+      proportion: 0.9  # Set to rate < 1 if training conditional generation
+      data_root: ./data/fMRI
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.fMRIDataset
+    params:
+      name: fMRI
+      proportion: 0.9  # rate
+      data_root: ./data/fMRI
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: test
+      style: separate
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 250
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/csdi/revenue-baseline-365.yaml

@@ -0,0 +1,82 @@
+model:
+  target: models.CSDI.tiffusion.Tiffusion
+  params:
+    seq_length: 365
+    feature_size: 3
+    n_layer_enc: 6
+    n_layer_dec: 4
+    d_model: 128  # 4 X 16
+    timesteps: 500 # diffusion timesteps
+    sampling_timesteps: 200
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 8
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+      # - classifier-based-sum-control
+      # - classifier-free-sum-control
+      # - range-wise-peak-control
+
+solver:
+  base_lr: 2.0e-3
+  max_epochs: 2230 # 11150
+  results_folder: ../../../data/CSDI/ckpt_baseline
+  gradient_accumulate_every: 2
+  save_cycle: 223 # 1115  # max_epochs // 5
+  ema:
+    decay: 0.995
+    update_interval: 10
+
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.65
+      patience: 200
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 100
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # Set to rate < 1 if training conditional generation
+      # data_root: ./Data/datasets/stock_data.csv
+      data_root: ../../../data/daily.csv
+      window: 365  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # rate
+      data_root: ../../../data/daily.csv
+      window: 365  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: test
+      style: separate
+      # distribution: geometric
+      distribution: uniform
+      missing_ratio: 0.5
+
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 100
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/csdi/sines.yaml

@@ -0,0 +1,73 @@
+model:
+  target: models.CSDI.tiffusion.Tiffusion
+  params:
+    seq_length: 24
+    feature_size: 5
+    n_layer_enc: 1
+    n_layer_dec: 2
+    d_model: 64  # 4 X 16
+    timesteps: 500
+    sampling_timesteps: 500
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+
+solver:
+  base_lr: 1.0e-3
+  max_epochs: 12000
+  results_folder: ../../../data/CSDI/Checkpoints_sine
+  gradient_accumulate_every: 2
+  save_cycle: 1200  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 3000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.sine_dataset.SineDataset
+    params:
+      num: 10000
+      dim: 5
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.sine_dataset.SineDataset
+    params:
+      num: 1000
+      dim: 5
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      style: separate
+      period: test
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 200
+
+  batch_size: 128
+  sample_size: 256
+  shuffle: True

config/energy.yaml

@@ -0,0 +1,74 @@
+model:
+  target: Models.interpretable_diffusion.gaussian_diffusion.Diffusion_TS
+  params:
+    seq_length: 24
+    feature_size: 28
+    n_layer_enc: 4
+    n_layer_dec: 3
+    d_model: 96  # 4 X 24
+    timesteps: 1000
+    sampling_timesteps: 1000
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 25000
+  results_folder: ./Checkpoints_energy
+  gradient_accumulate_every: 2
+  save_cycle: 2500  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 5000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.CustomDataset
+    params:
+      name: energy
+      proportion: 1.0  # Set to rate < 1 if training conditional generation
+      data_root: ./Data/datasets/energy_data.csv
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.CustomDataset
+    params:
+      name: energy
+      proportion: 0.9  # rate
+      data_root: ./Data/datasets/energy_data.csv
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: test
+      style: separate
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 250
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/etth.yaml

@@ -0,0 +1,74 @@
+model:
+  target: Models.interpretable_diffusion.gaussian_diffusion.Diffusion_TS
+  params:
+    seq_length: 24
+    feature_size: 7
+    n_layer_enc: 3
+    n_layer_dec: 2
+    d_model: 64  # 4 X 16
+    timesteps: 500
+    sampling_timesteps: 500
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 9000
+  results_folder: ./Checkpoints_etth
+  gradient_accumulate_every: 2
+  save_cycle: 1800  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 4000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.CustomDataset
+    params:
+      name: etth
+      proportion: 1.0  # Set to rate < 1 if training conditional generation
+      data_root: ./Data/datasets/ETTh.csv
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.CustomDataset
+    params:
+      name: etth
+      proportion: 0.9  # rate
+      data_root: ./Data/datasets/ETTh.csv
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: test
+      style: separate
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 200
+
+  batch_size: 128
+  sample_size: 256
+  shuffle: True

config/fmri.yaml

@@ -0,0 +1,74 @@
+model:
+  target: Models.interpretable_diffusion.gaussian_diffusion.Diffusion_TS
+  params:
+    seq_length: 24
+    feature_size: 50
+    n_layer_enc: 4
+    n_layer_dec: 4
+    d_model: 96  # 4 X 24
+    timesteps: 1000
+    sampling_timesteps: 1000
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 5
+    padding_size: 2
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 15000
+  results_folder: ./Checkpoints_fmri
+  gradient_accumulate_every: 2
+  save_cycle: 1500  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 3000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.fMRIDataset
+    params:
+      name: fMRI
+      proportion: 1.0  # Set to rate < 1 if training conditional generation
+      data_root: ./Data/datasets/fMRI
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.fMRIDataset
+    params:
+      name: fMRI
+      proportion: 0.9  # rate
+      data_root: ./Data/datasets/fMRI
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: test
+      style: separate
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 250
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/modified/192/energy.yaml

@@ -0,0 +1,74 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 192
+    feature_size: 28
+    n_layer_enc: 4
+    n_layer_dec: 3
+    d_model: 96  # 4 X 24
+    timesteps: 1000
+    sampling_timesteps: 1000
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 25000
+  results_folder: ../../../data/Checkpoints_energy
+  gradient_accumulate_every: 2
+  save_cycle: 2500  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 5000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.CustomDataset
+    params:
+      name: energy
+      proportion: 1.0  # Set to rate < 1 if training conditional generation
+      data_root: ./data/energy_data.csv
+      window: 192  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.CustomDataset
+    params:
+      name: energy
+      proportion: 0.9  # rate
+      data_root: ./data/energy_data.csv
+      window: 192  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: test
+      style: separate
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 250
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/modified/192/fmri.yaml

@@ -0,0 +1,74 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 192
+    feature_size: 50
+    n_layer_enc: 4
+    n_layer_dec: 4
+    d_model: 96  # 4 X 24
+    timesteps: 1000
+    sampling_timesteps: 1000
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 5
+    padding_size: 2
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 15000
+  results_folder: ../../../data/Checkpoints_fmri
+  gradient_accumulate_every: 2
+  save_cycle: 1500  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 3000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.fMRIDataset
+    params:
+      name: fMRI
+      proportion: 0.9  # Set to rate < 1 if training conditional generation
+      data_root: ./data/fMRI
+      window: 192  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.fMRIDataset
+    params:
+      name: fMRI
+      proportion: 0.9  # rate
+      data_root: ./data/fMRI
+      window: 192  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: test
+      style: separate
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 250
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/modified/192/revenue.yaml

@@ -0,0 +1,82 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 192
+    feature_size: 3
+    n_layer_enc: 6
+    n_layer_dec: 4
+    d_model: 128  # 4 X 16
+    timesteps: 500 # diffusion timesteps
+    sampling_timesteps: 200
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 8
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+      # - classifier-based-sum-control
+      # - classifier-free-sum-control
+      # - range-wise-peak-control
+
+solver:
+  base_lr: 2.0e-5
+  max_epochs: 2230 # 11150
+  results_folder: ../../../data/ckpt_baseline
+  gradient_accumulate_every: 2
+  save_cycle: 223 # 1115  # max_epochs // 5
+  ema:
+    decay: 0.995
+    update_interval: 10
+
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.65
+      patience: 200
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 100
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # Set to rate < 1 if training conditional generation
+      # data_root: ./Data/datasets/stock_data.csv
+      data_root: ../../../data/daily.csv
+      window: 192  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # rate
+      data_root: ../../../data/daily.csv
+      window: 192  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: test
+      style: separate
+      # distribution: geometric
+      distribution: uniform
+      missing_ratio: 0.5
+
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 100
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/modified/192/sines.yaml

@@ -0,0 +1,73 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 192
+    feature_size: 5
+    n_layer_enc: 1
+    n_layer_dec: 2
+    d_model: 64  # 4 X 16
+    timesteps: 500
+    sampling_timesteps: 500
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 12000
+  results_folder: ../../../data/Checkpoints_sine
+  gradient_accumulate_every: 2
+  save_cycle: 1200  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 3000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.sine_dataset.SineDataset
+    params:
+      num: 10000
+      dim: 5
+      window: 192  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.sine_dataset.SineDataset
+    params:
+      num: 1000
+      dim: 5
+      window: 192  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      style: separate
+      period: test
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 200
+
+  batch_size: 128
+  sample_size: 256
+  shuffle: True

config/modified/384/energy.yaml

@@ -0,0 +1,74 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 384
+    feature_size: 28
+    n_layer_enc: 4
+    n_layer_dec: 3
+    d_model: 96  # 4 X 24
+    timesteps: 1000
+    sampling_timesteps: 1000
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 25000
+  results_folder: ../../../data/Checkpoints_energy
+  gradient_accumulate_every: 2
+  save_cycle: 2500  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 5000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.CustomDataset
+    params:
+      name: energy
+      proportion: 1.0  # Set to rate < 1 if training conditional generation
+      data_root: ./data/energy_data.csv
+      window: 384  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.CustomDataset
+    params:
+      name: energy
+      proportion: 0.9  # rate
+      data_root: ./data/energy_data.csv
+      window: 384  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: test
+      style: separate
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 250
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/modified/384/fmri.yaml

@@ -0,0 +1,74 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 384
+    feature_size: 50
+    n_layer_enc: 4
+    n_layer_dec: 4
+    d_model: 96  # 4 X 24
+    timesteps: 1000
+    sampling_timesteps: 1000
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 5
+    padding_size: 2
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 15000
+  results_folder: ../../../data/Checkpoints_fmri
+  gradient_accumulate_every: 2
+  save_cycle: 1500  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 3000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.fMRIDataset
+    params:
+      name: fMRI
+      proportion: 0.9  # Set to rate < 1 if training conditional generation
+      data_root: ./data/fMRI
+      window: 384  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.fMRIDataset
+    params:
+      name: fMRI
+      proportion: 0.9  # rate
+      data_root: ./data/fMRI
+      window: 384  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: test
+      style: separate
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 250
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/modified/384/revenue.yaml

@@ -0,0 +1,82 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 384
+    feature_size: 3
+    n_layer_enc: 6
+    n_layer_dec: 4
+    d_model: 128  # 4 X 16
+    timesteps: 500 # diffusion timesteps
+    sampling_timesteps: 200
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 8
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+      # - classifier-based-sum-control
+      # - classifier-free-sum-control
+      # - range-wise-peak-control
+
+solver:
+  base_lr: 2.0e-5
+  max_epochs: 2230 # 11150
+  results_folder: ../../../data/ckpt_baseline
+  gradient_accumulate_every: 2
+  save_cycle: 223 # 1115  # max_epochs // 5
+  ema:
+    decay: 0.995
+    update_interval: 10
+
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.65
+      patience: 200
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 100
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # Set to rate < 1 if training conditional generation
+      # data_root: ./Data/datasets/stock_data.csv
+      data_root: ../../../data/daily.csv
+      window: 384  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # rate
+      data_root: ../../../data/daily.csv
+      window: 384  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: test
+      style: separate
+      # distribution: geometric
+      distribution: uniform
+      missing_ratio: 0.5
+
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 100
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/modified/384/sines.yaml

@@ -0,0 +1,73 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 384
+    feature_size: 5
+    n_layer_enc: 1
+    n_layer_dec: 2
+    d_model: 64  # 4 X 16
+    timesteps: 500
+    sampling_timesteps: 500
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 12000
+  results_folder: ../../../data/Checkpoints_sine
+  gradient_accumulate_every: 2
+  save_cycle: 1200  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 3000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.sine_dataset.SineDataset
+    params:
+      num: 10000
+      dim: 5
+      window: 384  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.sine_dataset.SineDataset
+    params:
+      num: 1000
+      dim: 5
+      window: 384  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      style: separate
+      period: test
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 200
+
+  batch_size: 128
+  sample_size: 256
+  shuffle: True

config/modified/96/energy.yaml

@@ -0,0 +1,74 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 96
+    feature_size: 28
+    n_layer_enc: 4
+    n_layer_dec: 3
+    d_model: 96  # 4 X 24
+    timesteps: 1000
+    sampling_timesteps: 1000
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 25000
+  results_folder: ../../../data/Checkpoints_energy
+  gradient_accumulate_every: 2
+  save_cycle: 2500  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 5000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.CustomDataset
+    params:
+      name: energy
+      proportion: 1.0  # Set to rate < 1 if training conditional generation
+      data_root: ./data/energy_data.csv
+      window: 96  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.CustomDataset
+    params:
+      name: energy
+      proportion: 0.9  # rate
+      data_root: ./data/energy_data.csv
+      window: 96  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: test
+      style: separate
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 250
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/modified/96/fmri.yaml

@@ -0,0 +1,74 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 96
+    feature_size: 50
+    n_layer_enc: 4
+    n_layer_dec: 4
+    d_model: 96  # 4 X 24
+    timesteps: 1000
+    sampling_timesteps: 1000
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 5
+    padding_size: 2
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 15000
+  results_folder: ../../../data/Checkpoints_fmri
+  gradient_accumulate_every: 2
+  save_cycle: 1500  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 3000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.fMRIDataset
+    params:
+      name: fMRI
+      proportion: 0.9  # Set to rate < 1 if training conditional generation
+      data_root: ./data/fMRI
+      window: 96  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.fMRIDataset
+    params:
+      name: fMRI
+      proportion: 0.9  # rate
+      data_root: ./data/fMRI
+      window: 96  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: test
+      style: separate
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 250
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/modified/96/revenue.yaml

@@ -0,0 +1,82 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 96
+    feature_size: 3
+    n_layer_enc: 6
+    n_layer_dec: 4
+    d_model: 128  # 4 X 16
+    timesteps: 500 # diffusion timesteps
+    sampling_timesteps: 200
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 8
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+      # - classifier-based-sum-control
+      # - classifier-free-sum-control
+      # - range-wise-peak-control
+
+solver:
+  base_lr: 2.0e-5
+  max_epochs: 2230 # 11150
+  results_folder: ../../../data/ckpt_baseline
+  gradient_accumulate_every: 2
+  save_cycle: 223 # 1115  # max_epochs // 5
+  ema:
+    decay: 0.995
+    update_interval: 10
+
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.65
+      patience: 200
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 100
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # Set to rate < 1 if training conditional generation
+      # data_root: ./Data/datasets/stock_data.csv
+      data_root: ../../../data/daily.csv
+      window: 96  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # rate
+      data_root: ../../../data/daily.csv
+      window: 96  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: test
+      style: separate
+      # distribution: geometric
+      distribution: uniform
+      missing_ratio: 0.5
+
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 100
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/modified/96/sines.yaml

@@ -0,0 +1,73 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 96
+    feature_size: 5
+    n_layer_enc: 1
+    n_layer_dec: 2
+    d_model: 64  # 4 X 16
+    timesteps: 500
+    sampling_timesteps: 500
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 12000
+  results_folder: ../../../data/Checkpoints_sine
+  gradient_accumulate_every: 2
+  save_cycle: 1200  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 3000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.sine_dataset.SineDataset
+    params:
+      num: 10000
+      dim: 5
+      window: 96  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.sine_dataset.SineDataset
+    params:
+      num: 1000
+      dim: 5
+      window: 96  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      style: separate
+      period: test
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 200
+
+  batch_size: 128
+  sample_size: 256
+  shuffle: True

config/modified/energy.yaml

@@ -0,0 +1,74 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 24
+    feature_size: 28
+    n_layer_enc: 4
+    n_layer_dec: 3
+    d_model: 96  # 4 X 24
+    timesteps: 1000
+    sampling_timesteps: 1000
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 25000
+  results_folder: ../../../data/Checkpoints_energy
+  gradient_accumulate_every: 2
+  save_cycle: 2500  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 5000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.CustomDataset
+    params:
+      name: energy
+      proportion: 1.0  # Set to rate < 1 if training conditional generation
+      data_root: ./data/energy_data.csv
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.CustomDataset
+    params:
+      name: energy
+      proportion: 0.9  # rate
+      data_root: ./data/energy_data.csv
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: test
+      style: separate
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 250
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/modified/fmri.yaml

@@ -0,0 +1,74 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 24
+    feature_size: 50
+    n_layer_enc: 4
+    n_layer_dec: 4
+    d_model: 96  # 4 X 24
+    timesteps: 1000
+    sampling_timesteps: 1000
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 5
+    padding_size: 2
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 15000
+  results_folder: ../../../data/Checkpoints_fmri
+  gradient_accumulate_every: 2
+  save_cycle: 1500  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 3000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.fMRIDataset
+    params:
+      name: fMRI
+      proportion: 0.9  # Set to rate < 1 if training conditional generation
+      data_root: ./data/fMRI
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.fMRIDataset
+    params:
+      name: fMRI
+      proportion: 0.9  # rate
+      data_root: ./data/fMRI
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: test
+      style: separate
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 250
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/modified/revenue-baseline-365.yaml

@@ -0,0 +1,82 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 365
+    feature_size: 3
+    n_layer_enc: 6
+    n_layer_dec: 4
+    d_model: 128  # 4 X 16
+    timesteps: 500 # diffusion timesteps
+    sampling_timesteps: 200
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 8
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+      # - classifier-based-sum-control
+      # - classifier-free-sum-control
+      # - range-wise-peak-control
+
+solver:
+  base_lr: 2.0e-5
+  max_epochs: 2230 # 11150
+  results_folder: ../../../data/ckpt_baseline
+  gradient_accumulate_every: 2
+  save_cycle: 223 # 1115  # max_epochs // 5
+  ema:
+    decay: 0.995
+    update_interval: 10
+
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.65
+      patience: 200
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 100
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # Set to rate < 1 if training conditional generation
+      # data_root: ./Data/datasets/stock_data.csv
+      data_root: ../../../data/daily.csv
+      window: 365  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # rate
+      data_root: ../../../data/daily.csv
+      window: 365  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: test
+      style: separate
+      # distribution: geometric
+      distribution: uniform
+      missing_ratio: 0.5
+
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 100
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/modified/revenue.yaml

@@ -0,0 +1,82 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 24
+    feature_size: 3
+    n_layer_enc: 6
+    n_layer_dec: 4
+    d_model: 128  # 4 X 16
+    timesteps: 500 # diffusion timesteps
+    sampling_timesteps: 200
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 8
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+      # - classifier-based-sum-control
+      # - classifier-free-sum-control
+      # - range-wise-peak-control
+
+solver:
+  base_lr: 2.0e-5
+  max_epochs: 2230 # 11150
+  results_folder: ../../../data/ckpt_baseline
+  gradient_accumulate_every: 2
+  save_cycle: 223 # 1115  # max_epochs // 5
+  ema:
+    decay: 0.995
+    update_interval: 10
+
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.65
+      patience: 200
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 100
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # Set to rate < 1 if training conditional generation
+      # data_root: ./Data/datasets/stock_data.csv
+      data_root: ../../../data/daily.csv
+      window: 365  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.RevenueDataset
+    params:
+      name: revenue
+      proportion: 0.8  # rate
+      data_root: ../../../data/daily.csv
+      window: 365  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 2024
+      period: test
+      style: separate
+      # distribution: geometric
+      distribution: uniform
+      missing_ratio: 0.5
+
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 100
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

config/modified/sines.yaml

@@ -0,0 +1,73 @@
+model:
+  target: models.Tiffusion.tiffusion.Tiffusion
+  params:
+    seq_length: 24
+    feature_size: 5
+    n_layer_enc: 1
+    n_layer_dec: 2
+    d_model: 64  # 4 X 16
+    timesteps: 500
+    sampling_timesteps: 500
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+    control_signal: []
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 12000
+  results_folder: ../../../data/Checkpoints_sine
+  gradient_accumulate_every: 2
+  save_cycle: 1200  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 3000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.sine_dataset.SineDataset
+    params:
+      num: 10000
+      dim: 5
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.sine_dataset.SineDataset
+    params:
+      num: 1000
+      dim: 5
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      style: separate
+      period: test
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 200
+
+  batch_size: 128
+  sample_size: 256
+  shuffle: True

config/mujoco.yaml

@@ -0,0 +1,72 @@
+model:
+  target: Models.interpretable_diffusion.gaussian_diffusion.Diffusion_TS
+  params:
+    seq_length: 24
+    feature_size: 14
+    n_layer_enc: 3
+    n_layer_dec: 2
+    d_model: 64  # 4 X 16
+    timesteps: 1000
+    sampling_timesteps: 1000
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 14000
+  results_folder: ./Checkpoints_mujoco
+  gradient_accumulate_every: 2
+  save_cycle: 1400  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 3000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.mujoco_dataset.MuJoCoDataset
+    params:
+      num: 10000
+      dim: 14
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.mujoco_dataset.MuJoCoDataset
+    params:
+      num: 1000
+      dim: 14
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      style: separate
+      period: test
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 250
+
+  batch_size: 128
+  sample_size: 256
+  shuffle: True

config/mujoco_sssd.yaml

@@ -0,0 +1,40 @@
+model:
+  target: Models.interpretable_diffusion.gaussian_diffusion.Diffusion_TS
+  params:
+    seq_length: 100
+    feature_size: 14
+    n_layer_enc: 3
+    n_layer_dec: 3
+    d_model: 64  # 4 X 16
+    timesteps: 500
+    sampling_timesteps: 500
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0
+    resid_pd: 0
+    kernel_size: 1
+    padding_size: 0
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 12000
+  results_folder: ./Checkpoints_mujoco_sssd
+  gradient_accumulate_every: 2
+  save_cycle: 1200  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 3000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False

config/sines.yaml

@@ -0,0 +1,72 @@
+model:
+  target: Models.interpretable_diffusion.gaussian_diffusion.Diffusion_TS
+  params:
+    seq_length: 24
+    feature_size: 5
+    n_layer_enc: 1
+    n_layer_dec: 2
+    d_model: 64  # 4 X 16
+    timesteps: 500
+    sampling_timesteps: 500
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 12000
+  results_folder: ./Checkpoints_sine
+  gradient_accumulate_every: 2
+  save_cycle: 1200  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 3000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.sine_dataset.SineDataset
+    params:
+      num: 10000
+      dim: 5
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.sine_dataset.SineDataset
+    params:
+      num: 1000
+      dim: 5
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      style: separate
+      period: test
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 200
+
+  batch_size: 128
+  sample_size: 256
+  shuffle: True

config/solar.yaml

@@ -0,0 +1,40 @@
+model:
+  target: Models.interpretable_diffusion.gaussian_diffusion.Diffusion_TS
+  params:
+    seq_length: 192
+    feature_size: 128
+    n_layer_enc: 4
+    n_layer_dec: 4
+    d_model: 96  # 4 X 24
+    timesteps: 500
+    sampling_timesteps: 500
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 1500
+  results_folder: ./Checkpoints_solar
+  gradient_accumulate_every: 2
+  save_cycle: 150  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 300
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 100 
+      verbose: False

config/solar_update.yaml

@@ -0,0 +1,40 @@
+model:
+  target: Models.interpretable_diffusion.gaussian_diffusion.Diffusion_TS
+  params:
+    seq_length: 192
+    feature_size: 137
+    n_layer_enc: 4
+    n_layer_dec: 4
+    d_model: 96  # 4 X 24
+    timesteps: 500
+    sampling_timesteps: 500
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.5
+    resid_pd: 0.5
+    kernel_size: 1
+    padding_size: 0
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 1000
+  results_folder: ./Checkpoints_solar_nips
+  gradient_accumulate_every: 2
+  save_cycle: 100  # max_epochs // 10
+  ema:
+    decay: 0.9
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 300
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 100 
+      verbose: False

config/stocks.yaml

@@ -0,0 +1,74 @@
+model:
+  target: Models.interpretable_diffusion.gaussian_diffusion.Diffusion_TS
+  params:
+    seq_length: 24
+    feature_size: 6
+    n_layer_enc: 2
+    n_layer_dec: 2
+    d_model: 64  # 4 X 16
+    timesteps: 500
+    sampling_timesteps: 500
+    loss_type: 'l1'
+    beta_schedule: 'cosine'
+    n_heads: 4
+    mlp_hidden_times: 4
+    attn_pd: 0.0
+    resid_pd: 0.0
+    kernel_size: 1
+    padding_size: 0
+
+solver:
+  base_lr: 1.0e-5
+  max_epochs: 10000
+  results_folder: ./Checkpoints_stock
+  gradient_accumulate_every: 2
+  save_cycle: 1000  # max_epochs // 10
+  ema:
+    decay: 0.995
+    update_interval: 10
+  
+  scheduler:
+    target: engine.lr_sch.ReduceLROnPlateauWithWarmup
+    params:
+      factor: 0.5
+      patience: 2000
+      min_lr: 1.0e-5
+      threshold: 1.0e-1
+      threshold_mode: rel
+      warmup_lr: 8.0e-4
+      warmup: 500 
+      verbose: False
+
+dataloader:
+  train_dataset:
+    target: utils.data_utils.real_datasets.CustomDataset
+    params:
+      name: stock
+      proportion: 1.0  # Set to rate < 1 if training conditional generation
+      data_root: ./Data/datasets/stock_data.csv
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: train
+
+  test_dataset:
+    target: utils.data_utils.real_datasets.CustomDataset
+    params:
+      name: stock
+      proportion: 0.9  # rate
+      data_root: ./Data/datasets/stock_data.csv
+      window: 24  # seq_length
+      save2npy: True
+      neg_one_to_one: True
+      seed: 123
+      period: test
+      style: separate
+      distribution: geometric
+    coefficient: 1.0e-2
+    step_size: 5.0e-2
+    sampling_steps: 200
+
+  batch_size: 64
+  sample_size: 256
+  shuffle: True

efficiency.py

@@ -0,0 +1,319 @@
+import os
+import torch
+import time
+
+os.environ["WANDB_ENABLED"] = "false"
+
+from engine.solver import Trainer
+from data.build_dataloader import build_dataloader
+from data.build_dataloader import build_dataloader_cond
+
+from utils.io_utils import load_yaml_config, instantiate_from_config
+import warnings
+
+warnings.simplefilter("ignore", UserWarning)
+
+import numpy as np
+
+import pickle
+from pathlib import Path
+
+
+def load_cached_results(cache_dir):
+    results = {"unconditional": None, "sum_controlled": {}, "anchor_controlled": {}}
+    for cache_file in cache_dir.glob("*.pkl"):
+        with open(cache_file, "rb") as f:
+            key = cache_file.stem
+            # if key=="unconditional":
+            #     continue
+            if key == "unconditional":
+                results["unconditional"] = pickle.load(f)
+            elif key.startswith("sum_"):
+                param = key[4:]  # Remove 'sum_' prefix
+                results["sum_controlled"][param] = pickle.load(f)
+            elif key.startswith("anchor_"):
+                param = key[7:]  # Remove 'anchor_' prefix
+                results["anchor_controlled"][param] = pickle.load(f)
+    return results
+
+
+def save_result(cache_dir, key, subkey, data):
+    return
+
+    if subkey:
+        filename = f"{key}_{subkey}.pkl"
+    else:
+        filename = f"{key}.pkl"
+    with open(cache_dir / filename, "wb") as f:
+        pickle.dump(data, f)
+
+
+class Arguments:
+    def __init__(self, config_path, gpu=0) -> None:
+        self.config_path = config_path
+        # self.config_path = "./config/control/revenue-baseline-sine.yaml"
+        self.save_dir = (
+            "../../../data/" + os.path.basename(self.config_path).split(".")[0]
+        )
+        self.gpu = gpu
+        os.makedirs(self.save_dir, exist_ok=True)
+
+        self.mode = "infill"
+        self.missing_ratio = 0.95
+        self.milestone = 10
+
+
+import argparse
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Controlled Sampling")
+    parser.add_argument(
+        "--config_path", type=str, default="./config/modified/energy.yaml"
+    )
+    parser.add_argument("--gpu", type=int, default=0)
+    return parser.parse_args()
+
+
+def run(run_args):
+
+    args = Arguments(run_args.config_path, run_args.gpu)
+    configs = load_yaml_config(args.config_path)
+    device = torch.device(f"cuda:{args.gpu}" if torch.cuda.is_available() else "cpu")
+    torch.cuda.set_device(args.gpu)
+
+    dl_info = build_dataloader(configs, args)
+    model = instantiate_from_config(configs["model"]).to(device)
+    trainer = Trainer(config=configs, args=args, model=model, dataloader=dl_info)
+    # args.milestone
+    trainer.load("10")
+    dataset = dl_info["dataset"]
+    test_dl_info = build_dataloader_cond(configs, args)
+    test_dataloader, test_dataset = test_dl_info["dataloader"], test_dl_info["dataset"]
+    coef = configs["dataloader"]["test_dataset"]["coefficient"]
+    stepsize = configs["dataloader"]["test_dataset"]["step_size"]
+    sampling_steps = configs["dataloader"]["test_dataset"]["sampling_steps"]
+    seq_length, feature_dim = test_dataset.window, test_dataset.var_num
+    dataset_name = os.path.basename(args.config_path).split(".")[0].split("-")[0]
+    mapper = {
+        "sines": "sines",
+        "revenue": "revenue",
+        "energy": "energy",
+        "fmri": "fMRI",
+    }
+    gap = seq_length // 5
+    if seq_length in [96, 192, 384]:
+        ori_data = np.load(
+            os.path.join(
+                "../../../data/train/",str(seq_length),
+                dataset_name,
+                "samples",
+                f'{mapper[dataset_name].replace("sines", "sine")}_norm_truth_{seq_length}_train.npy',
+            )
+        )
+        masks = np.load(
+            os.path.join(
+                "../../../data/train/",str(seq_length),
+                dataset_name,
+                "samples",
+                f'{mapper[dataset_name].replace("sines", "sine")}_masking_{seq_length}.npy',
+            )
+        )
+    else:
+        ori_data = np.load(
+            os.path.join(
+                "../../../data/train/",
+                dataset_name,
+                "samples",
+                f"{mapper[dataset_name]}_norm_truth_{seq_length}_train.npy",
+            )
+        )
+        masks = np.load(
+            os.path.join(
+                "../../../data/train/",
+                dataset_name,
+                "samples",
+                f"{mapper[dataset_name]}_masking_{seq_length}.npy",
+            )
+        )
+
+    sample_num, _, _ = masks.shape
+    # observed = ori_data[:sample_num] * masks
+    ori_data = ori_data[:sample_num]
+
+    sampling_size = min(1000, len(test_dataset), sample_num)
+    batch_size = 500
+    print(f"Sampling size: {sampling_size}, Batch size: {batch_size}")
+
+    ### Cache file path
+    cache_dir = Path(f"../../../data/cache/{dataset_name}_{seq_length}")
+    cache_dir.mkdir(exist_ok=True)
+    # results = load_cached_results(cache_dir)
+    results = {"unconditional": None, "sum_controlled": {}, "anchor_controlled": {}}
+
+    def measure_inference_time(func, *args, **kwargs):
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        return result, (end_time - start_time)
+
+    timing_results = {}
+
+    ### Unconditional sampling
+    if results["unconditional"] is None:
+        print("Generating unconditional data...")
+        results["unconditional"], timing = measure_inference_time(
+            trainer.control_sample,
+            num=sampling_size,
+            size_every=batch_size,
+            shape=[seq_length, feature_dim],
+            model_kwargs={
+                "gradient_control_signal": {},
+                "coef": coef,
+                "learning_rate": stepsize,
+            },
+        )
+        timing_results["unconditional"] = timing / sampling_size
+        save_result(cache_dir, "unconditional", "", results["unconditional"])
+
+    ### Different AUC values
+    auc_weights = [10]
+    auc_values = [-100, 20, 50, 150]  # -200, -150, -100, -50, 0, 20, 30, 50, 100, 150
+
+    for auc in auc_values:
+        for weight in auc_weights:
+            key = f"auc_{auc}_weight_{weight}"
+            if key not in results["sum_controlled"]:
+                print(f"Generating sum controlled data - AUC: {auc}, Weight: {weight}")
+                results["sum_controlled"][key], timing = measure_inference_time(
+                    trainer.control_sample,
+                    num=sampling_size,
+                    size_every=batch_size,
+                    shape=[seq_length, feature_dim],
+                    model_kwargs={
+                        "gradient_control_signal": {"auc": auc, "auc_weight": weight},
+                        "coef": coef,
+                        "learning_rate": stepsize,
+                    },
+                )
+                timing_results[f"sum_controlled_{key}"] = timing / sampling_size
+                save_result(cache_dir, "sum", key, results["sum_controlled"][key])
+
+    ### Different AUC weights
+    auc_weights = [1, 10, 50, 100]
+    auc_values = [-100]
+    for auc in auc_values:
+        for weight in auc_weights:
+            key = f"auc_{auc}_weight_{weight}"
+            if key not in results["sum_controlled"]:
+                print(f"Generating sum controlled data - AUC: {auc}, Weight: {weight}")
+                results["sum_controlled"][key], timing = measure_inference_time(
+                    trainer.control_sample,
+                    num=sampling_size,
+                    size_every=batch_size,
+                    shape=[seq_length, feature_dim],
+                    model_kwargs={
+                        "gradient_control_signal": {"auc": auc, "auc_weight": weight},
+                        "coef": coef,
+                        "learning_rate": stepsize,
+                    },
+                )
+                timing_results[f"sum_controlled_{key}"] = timing / (sampling_size)
+                save_result(cache_dir, "sum", key, results["sum_controlled"][key])
+
+
+    ### Different AUC segments
+    auc_weights = [10]
+    auc_values = [150]
+    auc_average = 10
+    auc_segments = ((gap, 2 * gap), (2 * gap, 3 * gap), (3 * gap, 4 * gap))
+    # for auc in auc_values:
+    #     for weight in auc_weights:
+    #         for segment in auc_segments:
+    auc = auc_values[0]
+    weight = auc_weights[0]
+    # segment = auc_segments[0]
+    for segment in auc_segments:
+        key = f"auc_{auc}_weight_{weight}_segment_{segment[0]}_{segment[1]}"
+        if key not in results["sum_controlled"]:
+            print(
+                f"Generating sum controlled data - AUC: {auc}, Weight: {weight}, Segment: {segment}"
+            )
+            results["sum_controlled"][key], timing = measure_inference_time(
+                trainer.control_sample,
+                num=sampling_size,
+                size_every=batch_size,
+                shape=[seq_length, feature_dim],
+                model_kwargs={
+                    "gradient_control_signal": {
+                        "auc": auc_average * (segment[1] - segment[0]), # / seq_length,
+                        "auc_weight": weight,
+                        "segment": [segment],
+                    },
+                    "coef": coef,
+                    "learning_rate": stepsize,
+                },
+            )
+            timing_results[f"sum_controlled_{key}"] = timing / sampling_size
+            save_result(cache_dir, "sum", key, results["sum_controlled"][key])
+
+    # Different anchors
+    anchor_values = [-0.8, 0.6, 1.0]
+    anchor_weights = [0.01, 0.01, 0.5, 1.0]
+    for peak in anchor_values:
+        for weight in anchor_weights:
+            key = f"peak_{peak}_weight_{weight}"
+            if key not in results["anchor_controlled"]:
+                mask = np.zeros((seq_length, feature_dim), dtype=np.float32)
+                mask[gap // 2 :: gap, 0] = weight
+                target = np.zeros((seq_length, feature_dim), dtype=np.float32)
+                target[gap // 2 :: gap, 0] = peak
+
+                print(f"Anchor controlled data - Peak: {peak}, Weight: {weight}")
+                results["anchor_controlled"][key], timing = measure_inference_time(
+                    trainer.control_sample,
+                    num=sampling_size,
+                    size_every=batch_size,
+                    shape=[seq_length, feature_dim],
+                    model_kwargs={
+                        "gradient_control_signal": {},  # "auc": -50, "auc_weight": 10.0},
+                        "coef": coef,
+                        "learning_rate": stepsize,
+                    },
+                    target=target,
+                    partial_mask=mask,
+                )
+                timing_results[f"anchor_controlled_{key}"] = timing / sampling_size
+                save_result(cache_dir, "anchor", key, results["anchor_controlled"][key])
+
+
+    ### Rerun Unconditional sampling
+    if results["unconditional"] is None:
+        print("Generating unconditional data...")
+        results["unconditional"], timing = measure_inference_time(
+            trainer.control_sample,
+            num=sampling_size,
+            size_every=batch_size,
+            shape=[seq_length, feature_dim],
+            model_kwargs={
+                "gradient_control_signal": {},
+                "coef": coef,
+                "learning_rate": stepsize,
+            },
+        )
+        timing_results["unconditional"] = timing / sampling_size
+        save_result(cache_dir, "unconditional", "", results["unconditional"])
+
+    # After all sampling is done, print timing results
+    print("\nAverage Inference Time per Sample:")
+    print("-" * 40)
+    for key, time_per_sample in timing_results.items():
+        print(f"{key}: {time_per_sample:.4f} seconds")
+
+    # return results, dataset_name, seq_length
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run(args)
+

engine/logger.py

@@ -0,0 +1,71 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import time
+import torch
+from utils.io_utils import write_args, save_config_to_yaml
+
+
+class Logger(object):
+    def __init__(self, args):
+        self.args = args
+        self.save_dir = args.save_dir
+
+        os.makedirs(self.save_dir, exist_ok=True)
+
+        # save the args and config
+        self.config_dir = os.path.join(self.save_dir, "configs")
+        os.makedirs(self.config_dir, exist_ok=True)
+        file_name = os.path.join(self.config_dir, "args.txt")
+        write_args(args, file_name)
+
+        log_dir = os.path.join(self.save_dir, "logs")
+        if not os.path.exists(log_dir):
+            os.makedirs(log_dir, exist_ok=True)
+        self.text_writer = open(os.path.join(log_dir, "log.txt"), "a")  # 'w')
+        if args.tensorboard:
+            self.log_info("using tensorboard")
+            self.tb_writer = torch.utils.tensorboard.SummaryWriter(
+                log_dir=log_dir
+            )  # tensorboard.SummaryWriter(log_dir=log_dir)
+        else:
+            self.tb_writer = None
+
+    def save_config(self, config):
+        save_config_to_yaml(config, os.path.join(self.config_dir, "config.yaml"))
+
+    def log_info(self, info, check_primary=True):
+        print(info)
+        info = str(info)
+        time_str = time.strftime("%Y-%m-%d-%H-%M")
+        info = "{}: {}".format(time_str, info)
+        if not info.endswith("\n"):
+            info += "\n"
+        self.text_writer.write(info)
+        self.text_writer.flush()
+
+    def add_scalar(self, **kargs):
+        """Log a scalar variable."""
+        if self.tb_writer is not None:
+            self.tb_writer.add_scalar(**kargs)
+
+    def add_scalars(self, **kargs):
+        """Log a scalar variable."""
+        if self.tb_writer is not None:
+            self.tb_writer.add_scalars(**kargs)
+
+    def add_image(self, **kargs):
+        """Log a scalar variable."""
+        if self.tb_writer is not None:
+            self.tb_writer.add_image(**kargs)
+
+    def add_images(self, **kargs):
+        """Log a scalar variable."""
+        if self.tb_writer is not None:
+            self.tb_writer.add_images(**kargs)
+
+    def close(self):
+        self.text_writer.close()
+        self.tb_writer.close()