diff --git a/scripts/training_data_generation/README_AI_RESTART_GENERATION.md b/scripts/training_data_generation/README_AI_RESTART_GENERATION.md
new file mode 100644
index 0000000..935e256
--- /dev/null
+++ b/scripts/training_data_generation/README_AI_RESTART_GENERATION.md
@@ -0,0 +1,146 @@
+# LandSim AI Model: Generating New Restart Files
+
+## Step 1: Prepare LandSim Inference Inputs
+
+### 1.1 Locate required input files
+
+You need the following files from your `tes_aoi_release` workflow and uELM run:
+
+- **Domain file**: `knox_domain.lnd.TES_SE.4km.1d.c*.nc` (from `tes_aoi_release/knox/domain_surfdata/`)
+- **Surfdata file**: `knox_surfdata.TES_SE.4km.1d.NLCD.c*.nc` (from `tes_aoi_release/knox/domain_surfdata/`)
+- **Forcing data**: Three subfolders in `tes_aoi_release/knox/forcing/`
+- **History files**: `*.elm.h0.*.nc` (from `${KMELM_ROOT}/e3sm_runs/uELM_knox_*/`)
+- **Restart file**: `*.elm.r.*.nc` (from `${KMELM_ROOT}/e3sm_runs/uELM_knox_*/`)
+
+### 1.2 Configure training data generation
+
+1. Navigate to the training data generation directory:
+
+   ```bash
+   cd /gpfs/.../LandSim/scripts/training_data_generation
+   ```
+
+2. Edit `config.py` to set the paths to your input files:
+
+   ```python
+   # Update these paths to match your environment
+   domain_file = "/path/to/knox_domain.lnd.TES_SE.4km.1d.c*.nc"
+   surfdata_file = "/path/to/knox_surfdata.TES_SE.4km.1d.NLCD.c*.nc"
+   forcing_dir = "/path/to/tes_aoi_release/knox/forcing"
+   history_files = "/path/to/kmELM/e3sm_runs/uELM_knox_*/history/*.elm.h0.*.nc"
+   restart_file = "/path/to/kmELM/e3sm_runs/uELM_knox_*/restart/*.elm.r.*.nc"
+   ```
+
+3. Review the README in `scripts/training_data_generation/` for detailed configuration options and command examples.
+
+### 1.3 Generate inference dataset
+
+Run the training data generation script to create the inference dataset:
+
+```bash
+cd /gpfs/.../LandSim/scripts/training_data_generation
+python enhanced_training_dataset.py --initial_only
+# or for enhanced dataset with final spinup data:
+# python enhanced_training_dataset.py --enhanced_dataset
+```
+
+This creates PKL files in the output directory (typically `output/initial_condition_dataset/` or `output/enhanced_dataset/`).
+
+## Step 2: Run LandSim Inference
+
+### 2.1 Activate LandSim environment
+
+```bash
+cd /gpfs/.../LandSim
+source .venv/bin/activate
+```
+
+### 2.2 Run inference
+
+From the LandSim repository root, run:
+
+```bash
+python scripts/run_inference_all.py \
+  --model /path/to/trained/model.pth \
+  --data-paths /gpfs/.../LandSim/scripts/training_data_generation/output/initial_condition_dataset \
+  --file-pattern "enhanced_monthly_training_data_batch_*.pkl" \
+  --output-dir /gpfs/.../LandSim/cnp_inference_knox_initial \
+  --use-training-config \
+  --strict-loading
+```
+
+**Parameters:**
+- `--model`: Path to your trained model file (`.pth`)
+- `--data-paths`: Directory containing the PKL files generated in Step 1
+- `--file-pattern`: Pattern matching your PKL files (adjust based on dataset type)
+- `--output-dir`: Directory where inference results will be saved
+- `--use-training-config`: Use the training configuration for normalization
+- `--strict-loading`: Enforce strict loading requirements
+
+**Note:** Adjust the `--file-pattern` based on your dataset:
+- For initial condition dataset: `"initial_condition_batch_*.pkl"` or `"enhanced_monthly_training_data_batch_*.pkl"`
+- For enhanced dataset with final spinup: `"enhanced_1_training_data_batch_*.pkl"`
+
+### 2.3 Verify inference output
+
+Check the output directory for generated files:
+
+```bash
+ls -lh /gpfs/.../LandSim/cnp_inference_knox_initial/cnp_predictions/
+```
+
+You should see:
+- `predictions_scalar.csv`
+- `pft_1d_predictions/` directory
+- `soil_2d_predictions/` directory
+- `predictions.pkl` file
+
+## Step 3: Convert AI Predictions to NetCDF
+
+Convert the AI predictions to NetCDF format for visualization and further processing:
+
+```bash
+cd /gpfs/.../LandSim/scripts
+python ai_predictions_to_netcdf.py \
+  --ai-predictions /gpfs/.../LandSim/cnp_inference_knox_initial/cnp_predictions \
+  --output /gpfs/.../LandSim/cnp_inference_knox_initial/ai_predictions_knox.nc \
+  --wrap-longitude
+```
+
+**Parameters:**
+- `--ai-predictions`: Directory containing the AI predictions from Step 2
+- `--output`: Output NetCDF file path
+- `--wrap-longitude`: Convert longitudes to -180–180 range for compatibility with plotting utilities
+
+The resulting NetCDF file (`ai_predictions_knox.nc`) can be used for:
+- Visualization with plotting scripts
+- Comparison with model outputs
+- Further analysis
+
+## Step 4: Create AI-Adjusted Restart File
+
+### 4.1 Update restart file with AI predictions
+
+Run the restart updater to create an AI-adjusted restart file. Only variables listed in the variable list file will be replaced with AI predictions; all other variables remain unchanged:
+
+```bash
+cd /gpfs/.../LandSim/scripts
+python ai_predictions_to_restart.py \
+  --ai-predictions /gpfs/.../LandSim/cnp_inference_knox_initial/ai_predictions_knox.nc \
+  --restart-file /gpfs/.../kmELM/e3sm_runs/uELM_knox_I1850CNPRDCTCBC/original_uELM_knox_I1850CNPRDCTCBC.elm.r.0021-01-01-00000.nc \
+  --variable-list /gpfs/.../LandSim/CNP_IO_updated14_xfer.txt \
+  --output /gpfs/.../LandSim/cnp_inference_knox_initial/uELM_knox_AIrestart.elm.r.0021-01-01-00000.nc
+```
+
+**Parameters:**
+- `--ai-predictions`: NetCDF file with AI predictions (from Step 3)
+- `--restart-file`: Original restart file from uELM accelerated spinup
+- `--variable-list`: Path to variable list file (e.g., `CNP_IO_updated14_xfer.txt`) that specifies which variables to update
+- `--output`: Path for the new AI-adjusted restart file
+
+
+### 4.2 Use AI-adjusted restart for final spinup
+
+The AI-adjusted restart file can now be used in the `tes_aoi_release` workflow for final spinup. See the `tes_aoi_release` README for instructions on using the restart file.
+
+
diff --git a/scripts/training_data_generation/config.py b/scripts/training_data_generation/config.py
index fac2849..73fac58 100644
--- a/scripts/training_data_generation/config.py
+++ b/scripts/training_data_generation/config.py
@@ -22,7 +22,8 @@
 output_dir = os.path.join(base_dir, 'output')
 
 # Processed forcing NetCDF files output directory
-forcing_netcdf_output_dir = os.path.join(output_dir, 'forcing_netcdf', 'TES_SE')
+# Updated to point to the actual location of processed forcing files
+forcing_netcdf_output_dir = '/gpfs/wolf2/cades/cli185/proj-shared/guzhuowei0407/Dataset_test/TVA/TVA_Forcing_netcdf'
 
 # Forcing PKL files output directory
 forcing_pkl_output_dir = os.path.join(output_dir, 'forcing_hourly_pkl')
@@ -31,7 +32,7 @@
 training_dataset_pkl_output_dir = os.path.join(output_dir, 'training_dataset_pkl')
 
 # CLM parameters NetCDF file path
-clm_params_nc_path = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/entire_domain/domain_surfdata/clm_params_c211124.nc'
+clm_params_nc_path = '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/clm_params_c211124.nc'
 
 # =============================================================================
 # INPUT FILES CONFIGURATION
@@ -39,25 +40,25 @@
 
 # Surface data files
 surface_data_files = [
-    '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/entire_domain/domain_surfdata/SEBOX1_surfdata.TES_SE.4km.1d.NLCD.c250202.nc'
+    '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/domain_surfdata/TVA_surfdata.TES_SE.4km.1d.NLCD.c241219.nc'
 ]
 
 # AD-SPINUP files (initial spinup)
 ad_spinup_history_files = [
-    '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/entire_domain/history_restart_files/uELM_SEBOX1_I1850CNPRDCTCBC.elm.h0.0021-01-01-00000.nc'
+    '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_adspinref.elm.h0.0021-01-01-00000.nc'
 ]
 
 ad_spinup_restart_files = [
-    '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/entire_domain/history_restart_files/uELM_SEBOX1_I1850CNPRDCTCBC.elm.r.0021-01-01-00000.nc'
+    '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_adspinref.elm.r.0021-01-01-00000.nc'
 ]
 
 # FINAL-SPINUP files (final spinup)
 final_spinup_history_files = [
-    #'/gpfs/wolf2/cades/cli185/proj-shared/wangd/kmELM/e3sm_runs/uELM_TVA_finalspinref/run/uELM_TVA_finalspinref.elm.h0.0781-01.nc'
+    '/gpfs/wolf2/cades/cli185/proj-shared/wangd/kmELM/e3sm_runs/uELM_TVA_finalspinref/run/uELM_TVA_finalspinref.elm.h0.0781-01.nc'
 ]
 
 final_spinup_restart_files = [
-    #'/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc'
+    '/gpfs/wolf2/cades/cli185/proj-shared/wangd/AI_data/TES_SE_dataset/TVA/history_restart_files/uELM_TVA_finalspinref.elm.r.0781-01-01-00000.nc'
 ]
 
 # Special P input NetCDF file
diff --git a/scripts/training_data_generation/python_scripts/CNP_IO_updated14_xfer.txt b/scripts/training_data_generation/python_scripts/CNP_IO_updated14_xfer.txt
index 8aa7de8..be12b23 100644
--- a/scripts/training_data_generation/python_scripts/CNP_IO_updated14_xfer.txt
+++ b/scripts/training_data_generation/python_scripts/CNP_IO_updated14_xfer.txt
@@ -59,4 +59,4 @@ SCALAR VARIABLES (1D - 4 variables):
 • soil3c_vr, soil3n_vr, soil3p_vr
 • soil4c_vr, soil4n_vr, soil4p_vr
 
-• labilep_vr , occlp_vr, primp_vr, secondp_vr
+• labilep_vr , occlp_vr, primp_vr, secondp_vr, solutionp_vr