fix: update notebook visualization

Author: mesakhcienet

examples/science/af3-slurm/examples/simple_ipynb_launcher/adm/PAE.png

@@ -285,7 +285,7 @@ elif [ "$JOB_TYPE" = "inference" ]; then
 
     # Add JAX cache path if set
     if [ -n "$JAX_COMPILATION_CACHE_PATH" ]; then
-        job_specific_python_opts+=" --jax_compilation_cache_path='${JAX_COMPILATION_CACHE_PATH}'"
+        job_specific_python_opts+=" --jax_compilation_cache_dir='${JAX_COMPILATION_CACHE_PATH}'"
         apptainer_binds+=("--bind" "${JAX_COMPILATION_CACHE_PATH}")
         echo "Info: JAX cache enabled at ${JAX_COMPILATION_CACHE_PATH}"
     else

examples/science/af3-slurm/examples/simple_ipynb_launcher/adm/datapipeline.png

@@ -205,7 +205,7 @@
     "datapipeline_partition_name = \"{{ default_datapipeline_partition_name }}\"\n",
     "\n",
     "# Assign `input_file` with datapipeline input file path\n",
-    "input_file= \"\" # Replace with your input file path (relative path from the jupyter notebook directory)\n"
+    "input_file= \"fold_input.json\" # Replace with your input file path (relative path from the jupyter notebook directory)\n"
    ]
   },
   {
@@ -259,8 +259,8 @@
     "# Datapipeline job could returns multiple output files in case the input file is a JSON file with multiple protein sequences\n",
     "# The output path is a list of dictionaries, each containing the path to the output file\n",
     "# The first output path is used as input for inference\n",
-    "print(\"[INFO] Extracting output paths from inference job response...\")\n",
-    "print(\"[INFO] Output count:\", len(datapipeline_submit_job_response[\"output_path_list\"]))\n",
+    "print(\"[INFO] Extracting output paths from datapipeline job response...\")\n",
+    "print(\"[INFO] Datapipeline output path count:\", len(datapipeline_submit_job_response[\"output_path_list\"]))\n",
     "inference_input_file = datapipeline_submit_job_response[\"output_path_list\"][0][\"output_path\"]\n",
     "\n",
     "# Assign `enable_unified_memory` with True or False\n",
@@ -328,7 +328,7 @@
     "# The output path is a list of dictionaries, each containing the path to the output file\n",
     "# In this example, we will use the first output path\n",
     "print(\"[INFO] Extracting output paths from inference job response...\")\n",
-    "print(\"[INFO] Output count:\", len(inference_submit_job_response[\"output_path_list\"]))\n",
+    "print(\"[INFO] Inference output path count:\", len(inference_submit_job_response[\"output_path_list\"]))\n",
     "inference_output_path_info = inference_submit_job_response[\"output_path_list\"][0]"
    ]
   },

examples/science/af3-slurm/examples/simple_ipynb_launcher/adm/inference.png

@@ -331,8 +331,6 @@ class AF3SlurmClient:
                 f"         └── Datapipeline output path : {datapipeline_output_file_path}")
 
         response["output_path_list"] = output_path_list
-        print(
-            f"[INFO] Datapipeline output path : {datapipeline_output_file_path}")
         return response
 
     def cancel_job(self, job_id: int) -> Optional[Dict]:

examples/science/af3-slurm/examples/simple_ipynb_launcher/adm/rest_api_launch_af3_job.sh.j2

@@ -26,9 +26,13 @@
 
 def plot_pae_matrix(pae: np.ndarray, token_chain_ids: str) -> None:
     df = pd.DataFrame(pae)
-    fig = px.imshow(df, color_continuous_scale='Viridis',
-                    labels={'color': 'PAE'},
-                    title='Predicted Aligned Error (PAE) Matrix')
+    fig = px.imshow(
+        df,
+        color_continuous_scale=[(0.0, 'green'), (1.0, 'white')],
+        labels={'color': 'PAE'},
+        title='Predicted Aligned Error (PAE) Matrix'
+    )
+
     # Draw chain boundaries if available
     if token_chain_ids:
         chain_boundaries = []
@@ -45,16 +49,17 @@ def plot_pae_matrix(pae: np.ndarray, token_chain_ids: str) -> None:
                 y0=0,
                 x1=boundary,
                 y1=len(token_chain_ids),
-                line=dict(color="red", width=1)
+                line=dict(color="black", width=1)
             )
             fig.add_shape(
                 type="line",
                 x0=0,
                 y0=boundary,
                 x1=len(token_chain_ids),
                 y1=boundary,
-                line=dict(color="red", width=1)
+                line=dict(color="black", width=1)
             )
+
     fig.update_layout(autosize=True)
     fig.show()
 

examples/science/af3-slurm/examples/simple_ipynb_launcher/adm/slurm-rest-api-notebook.ipynb.j2

@@ -254,7 +254,7 @@ elif [ "$JOB_TYPE" = "inference" ]; then
 
     # Add JAX cache path if set
     if [ -n "$JAX_COMPILATION_CACHE_PATH" ]; then
-        job_specific_python_opts+=" --jax_compilation_cache_path='${JAX_COMPILATION_CACHE_PATH}'"
+        job_specific_python_opts+=" --jax_compilation_cache_dir='${JAX_COMPILATION_CACHE_PATH}'"
         apptainer_binds+=("--bind" "${JAX_COMPILATION_CACHE_PATH}")
         echo "Info: JAX cache enabled at ${JAX_COMPILATION_CACHE_PATH}"
     else