update

app.py

@@ -43,8 +43,11 @@ class TimeSeriesEditor:
         # 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
-        
+        self.trending_controls = [
+            (200, 250, 0, self.function_parser.string_to_function("sin(2*pi*x)"), 0.05)
+            # 200,250,0,sin(2*pi*x),0.05
+        ]
+
     def format_value(self, value: float, feature_idx: int) -> str:
         """Format value with appropriate units and notation"""
         if self.show_normalized:
@@ -377,7 +380,7 @@ class TimeSeriesEditor:
                     peak_alpha: float,
                     auc_weight: float,
                     peak_weight: float,
-                    enable_trending: bool = False,
+                    enable_trending: bool = True,
                     enable_trending_with_diff: bool = False,
                     trending_params: str = ""
         ) -> Tuple[List[go.Figure], str, str, Dict]:
@@ -436,15 +439,16 @@ class TimeSeriesEditor:
         #         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
-        )
+        with torch.no_grad():
+            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
@@ -610,10 +614,10 @@ class TimeSeriesEditor:
 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]")
+        gr.Markdown("## Instruction: Scroll Down + Click [Update Figure] [~20s-30s] [Running on CPU...]")
 
         metrics_display = gr.JSON(label="Metrics", value={})
-        
+
         with gr.Row():
             with gr.Column(scale=1):
                 # with Tab():
@@ -642,103 +646,102 @@ def create_gradio_interface(editor: TimeSeriesEditor):
 
                 # 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)
+                # with gr.Accordion("Open for More Detail"):
+                with gr.Group():
+                    gr.Markdown("### Fixed 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 with Confidence")
+                    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.05"
+                    )
+
+                # 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)
                     
-                    # 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")
+                update_model_btn = gr.Button("Update Figure")
 
                 gr.Markdown("## Extend Edit", visible=False)
                 with gr.Tab("Range Shift", visible=False):
@@ -905,26 +908,26 @@ def create_gradio_interface(editor: TimeSeriesEditor):
                 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]
-        # )
+        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
 
@@ -1059,63 +1062,52 @@ class FunctionParser:
                 
             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"]
+
+    device = torch.device(f"cuda:0") if torch.cuda.is_available() else "cpu"
+    
+    from models.Tiffusion import tiffusion
+
+    model = tiffusion.Tiffusion(
+        seq_length=365,
+        feature_size=3,
+        n_layer_enc=6,
+        n_layer_dec=4,
+        d_model=128,
+        timesteps=500,
+        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=[]
+    ).to(device)
+
+    model.load_state_dict(torch.load("./weight/checkpoint-10.pt", map_location=device, weights_only=True)["model"])
+
+    coef = 1.0e-2
+    stepsize = 5.0e-2
+    sampling_steps = 100 # Adjustable between 100-500 for speed/accuracy tradeoff
+    seq_length = 365
     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 = TimeSeriesEditor(seq_length, feature_dim, model)
     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)

models/CSDI/tiffusion.py

@@ -33,7 +33,7 @@ def cosine_beta_schedule(timesteps, s=0.008):
     return torch.clip(betas, 0, 0.999)
 
 
-class Tiffusion(CSDI_base):
+class Tiffusion(nn.Module):
     def __init__(
         self,
         seq_length,
@@ -111,12 +111,9 @@ class Tiffusion(CSDI_base):
                 config_diff["beta_start"], config_diff["beta_end"], self.num_steps
             )
 
-        
         self.alpha_hat = 1 - self.beta
         self.alpha = np.cumprod(self.alpha_hat)
         self.alpha_torch = torch.tensor(self.alpha).float().to(self.device).unsqueeze(1).unsqueeze(1)
-        # self.beta = torch.from_numpy(self.beta).float().to(self.device)
-        # self.alpha = torch.from_numpy(self.alpha).float().to(self.device)
         
         self.emb_total_dim = self.emb_time_dim + self.emb_feature_dim
         if self.is_unconditional == False:
@@ -127,63 +124,8 @@ class Tiffusion(CSDI_base):
             num_embeddings=self.target_dim
             , embedding_dim=self.emb_feature_dim
         )
-        # self.model: Transformer = Transformer(
-        #     n_feat=feature_size,
-        #     n_channel=seq_length,
-        #     n_layer_enc=n_layer_enc,
-        #     n_layer_dec=n_layer_dec,
-        #     n_heads=n_heads,
-        #     attn_pdrop=attn_pd,
-        #     resid_pdrop=resid_pd,
-        #     mlp_hidden_times=mlp_hidden_times,
-        #     max_len=seq_length,
-        #     n_embd=d_model,
-        #     conv_params=[kernel_size, padding_size],
-        #     **kwargs,
-        # )
-        class Config:
-            def __init__(self, **kwargs):
-                self.__dict__.update(kwargs)
                 
-            # type: CSDI
-            # layers: 3
-            # channels: 64
-            # nheads: 8
-            # diffusion_embedding_dim: 128
-            # is_linear: False  # linear transformer
-
-            # beta_start: 0.0001
-            # beta_end: 0.5
-            # schedule: "quad"
-
-            # num_steps: 50
-
-            # # edit
-            # edit_steps: 50  # the number of steps to perform editing
-            # bootstrap_ratio: 0.5   # [0,1]
-
-            # is_attr_proj: False
-
-            # side:
-            # num_var: 1
-            # var_emb: 16
-            # time_emb: 128
-            
-            # attrs:
-            # attr_emb: 64
-        # config_diff["side_dim"] = self.emb_total_dim
         self.diffmodel = diff_CSDI(
-            # config=Config(
-            #     layers=3,
-            #     channels=64,
-            #     nheads=8,
-            #     diffusion_embedding_dim=128,
-            #     is_linear=False,
-            #     beta_start=0.0001,
-            #     beta_end=0.5,
-            #     schedule="quad",
-            #     num_steps=50,
-            # )
             {
                 "layers": 3,
                 "channels": 64,