20260203 meeting

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

# New: Integrate AutoFRK for spatial prediction
if enable_frk and loc is not None:
    temp = epsilon_theta.permute(1, 0, 2)
    #print(f"temp shape: {temp.shape}")
    epsilon_theta_fused = torch.zeros_like(epsilon_theta)
    #print(f"epsilon_theta_fused shape: {epsilon_theta_fused.shape}")
    for i in range(temp.shape[0]):
        success = False
        while not success:
            try:
                df = temp[i] + 1e-8 * torch.randn_like(temp[i])
                #print(f"df shape: {df.shape}")
                #print(f"loc shape: {loc.shape}")
                frk_model = AutoFRK(device=df.device, dtype=df.dtype, logger_level=30)
                frk_model.forward(data=df, loc=loc, requires_grad=True)
                pred_res = frk_model.predict(newloc=loc)
                frk_pred = pred_res['pred.value']
                epsilon_theta_fused[:, i, :] = frk_pred
                success = True  # successful and exit loop
            except Exception as e:
                LOGGER.warning(f"Failed to process record {i}. Retrying. Error: {e}")
elif enable_frk and loc is None:
    raise ValueError("Location data must be provided when enable_frk is True.")
else:
    epsilon_theta_fused = epsilon_theta

# Compute loss
if generate_only_missing:
    loss = loss_function(epsilon_theta_fused[loss_mask], noise[loss_mask])
else:
    loss = loss_function(epsilon_theta_fused, noise)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

if self.enable_spatial_prediction:

    # new feature for this chunk
    sp_mean = sssd_inference.mean(dim=0, keepdim=True)
    sp_std = sssd_inference.std(dim=0, unbiased=False, keepdim=True)
    sssd_inference = (sssd_inference - sp_mean) / sp_std

    autoFRK_inference = self._autoFRK_generate(
        sssd_inference = sssd_inference,
        with_known_loc = True
    )

    # new feature for this chunk
    autoFRK_inference = autoFRK_inference * sp_std + sp_mean

    results = {'imputation': autoFRK_inference.detach().cpu().numpy()}
else:
    results = {'imputation': sssd_inference.detach().cpu().numpy()}
self._save_data(results, 0)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

wavenet:
  # WaveNet model parameters
  input_channels: 8  # Number of input channels
  output_channels: 8  # Number of output channels
  residual_layers: 20  # Number of residual layers
  residual_channels: 64  # Number of channels in residual blocks
  skip_channels: 64  # Number of channels in skip connections

  # Diffusion step embedding dimensions
  diffusion_step_embed_dim_input: 128  # Input dimension
  diffusion_step_embed_dim_hidden: 256  # Middle dimension
  diffusion_step_embed_dim_output: 256  # Output dimension

  # Structured State Spaces sequence model (S4) configurations
  s4_max_sequence_length: 986  # Maximum sequence length
  s4_state_dim: 128  # State dimension
  s4_dropout: 0.1  # Dropout rate
  s4_bidirectional: true  # Whether to use bidirectional layers
  s4_use_layer_norm: true  # Whether to use layer normalization

diffusion:
  # Diffusion model parameters
  T: 500  # Number of diffusion steps
  beta_0: 0.0008  # Initial beta value
  beta_T: 0.08  # Final beta value

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

# Training configuration
batch_size: 50  # Batch size
output_directory: "/home/u6025091/SSSD_CP/results/Weather2k-fast-rectangular"  # Output directory for checkpoints and logs
ckpt_iter: "max"  # Checkpoint mode (max or min)
iters_per_ckpt: 500  # Checkpoint frequency (number of epochs)
iters_per_logging: 200  # Log frequency (number of iterations)
n_iters: 38000  # Maximum number of iterations
learning_rate: 0.0005  # Learning rate

# Additional training settings
only_generate_missing: true  # Generate missing values only
use_model: 2  # Model to use for training
masking: "forecast"  # Masking strategy for missing values
missing_k: 33  # Number of missing values

# Data paths
data:
  train_path: "/home/u6025091/SSSD_CP/datasets/Weather2k/data_train_known_real.npy"  # Path to training data

# autoFRK config
enable_spatial_training: false  # Enable spatial training step
location_path: "/home/u6025091/SSSD_CP/datasets/Weather2k/stations_known_locations.npy"  # Path to known locations
AFRK_method: "fast"  # waitting for adding into the model
AFRK_tps_method: "rectangular"  # waitting for adding into the model

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

# Inference configuration
batch_size: 50  # Batch size for inference
output_directory: "/home/u6025091/SSSD_CP/results/Weather2k-fast-rectangular/inference"  # Output directory for inference results
ckpt_path: "/home/u6025091/SSSD_CP/results/Weather2k-fast-rectangular"  # Path to checkpoint for inference
trials: 1 # Replications

# Additional training settings
only_generate_missing: true  # Generate missing values only
use_model: 2  # Model to use for training
masking: "forecast"  # Masking strategy for missing values
missing_k: 33  # Number of missing values

# Data paths
data:
  test_path: "/home/u6025091/SSSD_CP/datasets/Weather2k/data_test_missing.npy"  # Path to test data

# autoFRK config
enable_spatial_inference: true  # Enable spatial prediction step
known_location_path: "/home/u6025091/SSSD_CP/datasets/Weather2k/stations_known_locations.npy"  # Path to known locations
unknown_location_path: "/home/u6025091/SSSD_CP/datasets/Weather2k/stations_unknown_locations.npy"  # Path to unknown locations
AFRK_method: "fast"
AFRK_tps_method: "rectangular"