Update app.py

app.py

@@ -1,207 +1,514 @@
-import csv
-import sys
+import os
+import math
+import tempfile
+import gradio as gr
+import cudaq
 import numpy as np
+import cupy as cp
+from pathlib import Path
 import plotly.graph_objects as go
-import gradio as gr
+import plotly.io as pio
+import imageio # Keep for potential future use, but GIF generation removed
 from scipy.spatial import Delaunay
-import traceback
 
-# --- FIX for CSV Field Size Limit ---
-# Increase the CSV field size limit to handle large Plotly JSON objects that are
-# generated when Gradio's Examples component tries to log the output.
-max_int = sys.maxsize
-while True:
-    try:
-        csv.field_size_limit(max_int)
-        break
-    except OverflowError:
-        max_int = int(max_int / 10)
-# -----------------------------------------
-
-
-def solve_and_plot_interactive(Lx: float,
-                               Ly: float,
-                               t_max: float,
-                               M: int,
-                               Gamma: float = 0.1,
-                               Nx: int = 50,
-                               Ny: int = 50,
-                               initial: str = "gaussian",
-                               bc: str = "dirichlet"):
+# Set Plotly engine for image export
+try:
+    pio.kaleido.scope.mathjax = None
+except AttributeError:
+    pass
+
+def simulate_qlbm_and_animate(num_reg_qubits: int, T: int, distribution_type: str, ux_input: float, uy_input: float, velocity_field_type: str):
     """
-    Solves the 2D heat equation and returns an interactive Plotly figure with a fixed axis range.
+    Simulates a 2D advection-diffusion problem using a Quantum Lattice Boltzmann Method (QLBM)
+    and generates an interactive Plotly figure with a slider for selected time steps.
     """
-    # --- 1. Simulation Setup (No changes here) ---
-    x = np.linspace(0, Lx, Nx)
-    y = np.linspace(0, Ly, Ny)
-    dx, dy = x[1] - x[0], y[1] - y[0]
-    dt = 0.9 / (2 * Gamma * (1/dx**2 + 1/dy**2))
-    Nt = int(np.ceil(t_max / dt)) + 1
-    rx, ry = Gamma * dt / dx**2, Gamma * dt / dy**2
-    X, Y = np.meshgrid(x, y, indexing='ij')
-
-    u = np.zeros((Nx, Ny))
-    if initial == "gaussian":
-        u = np.exp(-(((X - Lx/2)**2 + (Y - Ly/2)**2) / (2*(Lx/10)**2)))
-    elif initial == "random":
-        u = np.random.rand(Nx, Ny)
-    elif initial == "sinusoidal":
-        kx, ky = 2 * np.pi / Lx, 2 * np.pi / Ly
-        u = np.sin(kx * X) * np.sin(ky * Y)
-    elif initial == "step":
-        u = np.where((X < Lx/2) & (Y < Ly/2), 1.0, 0.0)
-
-    # --- 2. Solve the Heat Equation (No changes here) ---
-    time_indices = np.linspace(0, Nt - 1, M, dtype=int)
-    U_slider = np.zeros((M, Nx, Ny))
-    store_idx = 0
-    if 0 in time_indices:
-        U_slider[store_idx] = u.copy()
-        store_idx += 1
-
-    for n in range(1, Nt):
-        un = u.copy()
-        u[1:-1, 1:-1] = (
-            un[1:-1, 1:-1]
-            + rx * (un[2:, 1:-1] - 2 * un[1:-1, 1:-1] + un[:-2, 1:-1])
-            + ry * (un[1:-1, 2:] - 2 * un[1:-1, 1:-1] + un[1:-1, :-2])
-        )
-        if bc == "dirichlet":
-            u[0, :] = u[-1, :] = u[:, 0] = u[:, -1] = 0.0
-        elif bc == "neumann":
-            u[0, :], u[-1, :], u[:, 0], u[:, -1] = u[1, :], u[-2, :], u[:, 1], u[:, -2]
-        elif bc == "periodic":
-            u[0, :], u[-1, :], u[:, 0], u[:, -1] = un[-2, :], un[1, :], un[:, -2], un[:, 1]
+    # GIF related variables removed as GIF generation is no longer needed
+    # video_length = T
+    # simulation_fps = 0.1
+    # frames = int(simulation_fps * video_length)
+    # if frames == 0:
+    #     frames = 1
+
+    num_anc = 3
+    num_qubits_total = 2 * num_reg_qubits + num_anc
+    current_N = 2**num_reg_qubits
+    N_tot_state_vector = 2**num_qubits_total
+    num_ranks = 1
+    rank = 0
+    N_sub_per_rank = int(N_tot_state_vector // num_ranks)
+
+    # timesteps_per_frame logic removed as it was for GIF
+    # timesteps_per_frame = 1
+    # if frames < T and frames > 0:
+    #     timesteps_per_frame = int(T / frames)
+    # if timesteps_per_frame == 0:
+    #     timesteps_per_frame = 1
+
+    # Initial state setup
+    if distribution_type == "Sine Wave (Original)":
+        selected_initial_state_function_raw = lambda x, y, N_val_func: \
+            np.sin(x * 2 * np.pi / N_val_func) * (1 - 0.5 * x / N_val_func) * \
+            np.sin(y * 4 * np.pi / N_val_func) * (1 - 0.5 * y / N_val_func) + 1
+    elif distribution_type == "Gaussian":
+        selected_initial_state_function_raw = lambda x, y, N_val_func: \
+            np.exp(-((x - N_val_func / 2)**2 / (2 * (N_val_func / 5)**2) +
+                     (y - N_val_func / 2)**2 / (2 * (N_val_func / 5)**2))) * 1.8 + 0.2
+    elif distribution_type == "Random":
+        selected_initial_state_function_raw = lambda x, y, N_val_func: \
+            np.random.rand(N_val_func, N_val_func) * 1.5 + 0.2 if isinstance(x, int) else \
+            np.random.rand(x.shape, x.shape[3]) * 1.5 + 0.2
+    else:
+        print(f"Warning: Unknown distribution type '{distribution_type}'. Defaulting to Sine Wave.")
+        selected_initial_state_function_raw = lambda x, y, N_val_func: \
+            np.sin(x * 2 * np.pi / N_val_func) * (1 - 0.5 * x / N_val_func) * \
+            np.sin(y * 4 * np.pi / N_val_func) * (1 - 0.5 * y / N_val_func) + 1
+
+    initial_state_func_eval = lambda x_coords, y_coords: \
+        selected_initial_state_function_raw(x_coords, y_coords, current_N) * \
+        (y_coords < current_N).astype(int)
+
+    with tempfile.TemporaryDirectory() as tmp_npy_dir:
+        intermediate_folder_path = Path(tmp_npy_dir)
+
+        cudaq.set_target('nvidia', option='fp64')
+
+        @cudaq.kernel
+        def alloc_kernel(num_qubits_alloc: int):
+            qubits = cudaq.qvector(num_qubits_alloc)
+
+        from cupy.cuda.memory import MemoryPointer, UnownedMemory
+
+        def to_cupy_array(state):
+            tensor = state.getTensor()
+            pDevice = tensor.data()
+            sizeByte = tensor.get_num_elements() * tensor.get_element_size()
+            mem = UnownedMemory(pDevice, sizeByte, owner=state)
+            memptr_obj = MemoryPointer(mem, 0)
+            cupy_array_val = cp.ndarray(tensor.get_num_elements(),
+                                        dtype=cp.complex128,
+                                        memptr=memptr_obj)
+            return cupy_array_val
+
+        class QLBMAdvecDiffD2Q5_new:
+            def __init__(self, ux=0.2, uy=0.15) -> None:
+                self.dim = 2
+                self.ndir = 5
+                self.nq_dir = math.ceil(np.log2(self.ndir))
+                self.dirs =
+                for dir_int in range(self.ndir):
+                    dir_bin = f"{dir_int:b}".zfill(self.nq_dir)
+                    self.dirs.append(dir_bin)
+                self.e_unitvec = np.array([0, 1, -1, 1, -1])
+                self.wts = np.array([2/6, 1/6, 1/6, 1/6, 1/6])
+                self.cs = 1 / np.sqrt(3)
+                self.ux = ux
+                self.uy = uy
+                self.u = np.array([0, self.ux, self.ux, self.uy, self.uy])
+                self.wtcoeffs = np.multiply(self.wts, 1 + self.e_unitvec * self.u / self.cs**2)
+                self.create_circuit()
+
+            def create_circuit(self):
+                v = np.pad(self.wtcoeffs, (0, 2**num_anc - self.ndir))
+                v = v**0.5
+                v += 1 # Original line was v += 1, not v += 1
+                v = v / np.linalg.norm(v)
+                U_prep = 2 * np.outer(v, v) - np.eye(len(v))
+                cudaq.register_operation("prep_op", U_prep)
+
+                def collisionOp(dirs_list):
+                    dirs_i_list_val =
+                    for dir_str in dirs_list:
+                        dirs_i = [(int(c)) for c in dir_str]
+                        dirs_i_list_val += dirs_i[::-1]
+                    return dirs_i_list_val
+
+                self.dirs_i_list = collisionOp(self.dirs)
+
+                @cudaq.kernel
+                def rshift(q: cudaq.qview, n: int):
+                    for i in range(n):
+                        if i == n - 1:
+                            x(q[n - 1 - i])
+                        elif i == n - 2:
+                            x.ctrl(q[n - 1 - (i + 1)], q[n - 1 - i])
+                        else:
+                            x.ctrl(q[0:n - 1 - i], q[n - 1 - i])
+
+                @cudaq.kernel
+                def lshift(q: cudaq.qview, n: int):
+                    for i in range(n):
+                        if i == 0:
+                            x(q) # Corrected from x(q)
+                        elif i == 1:
+                            x.ctrl(q, q[3])
+                        else:
+                            x.ctrl(q[0:i], q[i])
+
+                @cudaq.kernel
+                def d2q5_tstep(q: cudaq.qview, nqx: int, nqy: int, nq_dir_val: int, dirs_i_val: list[int]):
+                    qx = q[0:nqx]
+                    qy = q[nqx:nqx + nqy]
+                    qdir = q[nqx + nqy:nqx + nqy + nq_dir_val]
+                    
+                    idx_lqx = 2
+                    b_list = dirs_i_val[idx_lqx * nq_dir_val:(idx_lqx + 1) * nq_dir_val]
+                    for j in range(nq_dir_val):
+                        if b_list[j] == 0: x(qdir[j])
+                    cudaq.control(lshift, qdir, qx, nqx)
+                    for j in range(nq_dir_val):
+                        if b_list[j] == 0: x(qdir[j])
+
+                    idx_rqx = 1
+                    b_list = dirs_i_val[idx_rqx * nq_dir_val:(idx_rqx + 1) * nq_dir_val]
+                    for j in range(nq_dir_val):
+                        if b_list[j] == 0: x(qdir[j])
+                    cudaq.control(rshift, qdir, qx, nqx)
+                    for j in range(nq_dir_val):
+                        if b_list[j] == 0: x(qdir[j])
+
+                    idx_lqy = 4
+                    b_list = dirs_i_val[idx_lqy * nq_dir_val:(idx_lqy + 1) * nq_dir_val]
+                    for j in range(nq_dir_val):
+                        if b_list[j] == 0: x(qdir[j])
+                    cudaq.control(lshift, qdir, qy, nqy)
+                    for j in range(nq_dir_val):
+                        if b_list[j] == 0: x(qdir[j])
+
+                    idx_rqy = 3
+                    b_list = dirs_i_val[idx_rqy * nq_dir_val:(idx_rqy + 1) * nq_dir_val]
+                    for j in range(nq_dir_val):
+                        if b_list[j] == 0: x(qdir[j])
+                    cudaq.control(rshift, qdir, qy, nqy)
+                    for j in range(nq_dir_val):
+                        if b_list[j] == 0: x(qdir[j])
+
+                @cudaq.kernel
+                def d2q5_tstep_wrapper(state_arg: cudaq.State, nqx: int, nqy: int, nq_dir_val: int, dirs_i_val: list[int]):
+                    q = cudaq.qvector(state_arg)
+                    qdir = q[nqx + nqy:nqx + nqy + nq_dir_val]
+                    prep_op(qdir[4], qdir[3], qdir) # Corrected from qdir
+                    d2q5_tstep(q, nqx, nqy, nq_dir_val, dirs_i_val)
+                    prep_op(qdir[4], qdir[3], qdir) # Corrected from qdir
+
+                @cudaq.kernel
+                def d2q5_tstep_wrapper_hadamard(vec_arg: list[complex], nqx: int, nqy: int, nq_dir_val: int, dirs_i_val: list[int]):
+                    q = cudaq.qvector(vec_arg)
+                    qdir = q[nqx + nqy:nqx + nqy + nq_dir_val]
+                    qy = q[nqx:nqx + nqy]
+                    prep_op(qdir[4], qdir[3], qdir) # Corrected from qdir
+                    d2q5_tstep(q, nqx, nqy, nq_dir_val, dirs_i_val)
+                    prep_op(qdir[4], qdir[3], qdir) # Corrected from qdir
+                    for i in range(nqy):
+                        h(qy[i])
+
+                def run_timestep_func(vec_arg, hadamard=False):
+                    if hadamard:
+                        result = cudaq.get_state(d2q5_tstep_wrapper_hadamard, vec_arg, num_reg_qubits, num_reg_qubits, self.nq_dir, self.dirs_i_list)
+                    else:
+                        result = cudaq.get_state(d2q5_tstep_wrapper, vec_arg, num_reg_qubits, num_reg_qubits, self.nq_dir, self.dirs_i_list)
+                    num_nonzero_ranks = num_ranks / (2**num_anc)
+                    rank_slice_cupy = to_cupy_array(result)
+                    if rank >= num_nonzero_ranks and num_nonzero_ranks > 0:
+                        sub_sv_zeros = np.zeros(N_sub_per_rank, dtype=np.complex128)
+                        cp.cuda.runtime.memcpy(rank_slice_cupy.data.ptr, sub_sv_zeros.ctypes.data, sub_sv_zeros.nbytes, cp.cuda.runtime.memcpyHostToDevice)
+                    if rank == 0 and num_nonzero_ranks < 1 and N_sub_per_rank > 0:
+                        limit_idx = int(N_tot_state_vector / (2**num_anc))
+                        if limit_idx < rank_slice_cupy.size:
+                            rank_slice_cupy[limit_idx:] = 0
+                    return result
+                self.run_timestep = run_timestep_func
+
+            def write_state(self, state_to_write, t_step_str_val):
+                rank_slice_cupy = to_cupy_array(state_to_write)
+                num_nonzero_ranks = num_ranks / (2**num_anc)
+                if rank < num_nonzero_ranks or (rank == 0 and num_nonzero_ranks <= 0):
+                    save_path = intermediate_folder_path / f"{t_step_str_val}_{rank}.npy"
+                    with open(save_path, 'wb') as f:
+                        arr_to_save = None
+                        data_limit = N_sub_per_rank
+                        if num_nonzero_ranks < 1 and rank == 0:
+                            data_limit = int(N_tot_state_vector / (2**num_anc))
+                        if data_limit > 0:
+                            relevant_part_cupy = cp.real(rank_slice_cupy[:data_limit])
+                        else:
+                            relevant_part_cupy = cp.array(, dtype=cp.float64)
+                        if relevant_part_cupy.size >= current_N * current_N:
+                            arr_flat = relevant_part_cupy[:current_N * current_N]
+                            if downsampling_factor > 1 and current_N > 0:
+                                arr_reshaped = arr_flat.reshape((current_N, current_N))
+                                arr_downsampled = arr_reshaped[::downsampling_factor, ::downsampling_factor]
+                                arr_to_save = arr_downsampled.flatten()
+                            else:
+                                arr_to_save = arr_flat
+                        elif relevant_part_cupy.size > 0:
+                            if downsampling_factor > 1:
+                                arr_to_save = relevant_part_cupy[::downsampling_factor]
+                            else:
+                                arr_to_save = relevant_part_cupy
+                        if arr_to_save is not None and arr_to_save.size > 0:
+                            np.save(f, arr_to_save.get() if isinstance(arr_to_save, cp.ndarray) else arr_to_save)
+
+            def run_evolution(self, initial_state_arg, total_timesteps, observable=False, timesteps_to_save=None):
+                current_state_val = initial_state_arg
+                for t_iter in range(total_timesteps):
+                    next_state_val = None
+                    if t_iter == total_timesteps - 1 and observable:
+                        next_state_val = self.run_timestep(current_state_val, True)
+                        self.write_state(next_state_val, str(t_iter)) # Save final state
+                    else:
+                        next_state_val = self.run_timestep(current_state_val)
+                        # Save data only for specific intervals for the slider
+                        if timesteps_to_save and t_iter in timesteps_to_save:
+                            self.write_state(next_state_val, str(t_iter))
+                        if rank == 0 and t_iter % 10 == 0: # Print progress less frequently
+                            print(f"Timestep: {t_iter}/{total_timesteps}")
+                    cp.get_default_memory_pool().free_all_blocks()
+                    current_state_val = next_state_val
+                if rank == 0:
+                    print(f"Timestep: {total_timesteps}/{total_timesteps} (Evolution complete)")
+                cp.get_default_memory_pool().free_all_blocks()
+                self.final_state = current_state_val
+
+        downsampling_factor = 2**5
+        if current_N == 0:
+            print("Error: current_N is zero. num_reg_qubits likely too small.")
+            return None, None
+        if current_N < downsampling_factor:
+            downsampling_factor = current_N if current_N > 0 else 1
+
+        qlbm_obj = QLBMAdvecDiffD2Q5_new(ux=ux_input, uy=uy_input)
+        initial_state_val = cudaq.get_state(alloc_kernel, num_qubits_total)
         
-        if n in time_indices and store_idx < M:
-            U_slider[store_idx] = u.copy()
-            store_idx += 1
-            
-    # --- 3. Create a Plotly Figure for Web ---
-
-    # --- FIX: Calculate the global min/max temperature for a fixed Z-axis range ---
-    z_min = U_slider.min()
-    z_max = U_slider.max()
-    # Add a tiny buffer if the data is completely flat to avoid a display error
-    if z_max == z_min:
-        z_max += 1e-9
-    # --------------------------------------------------------------------------------
+        xv_init = np.arange(current_N)
+        yv_init = np.arange(current_N)
+        initial_grid_2d_X, initial_grid_2d_Y = np.meshgrid(xv_init, yv_init)
+
+        if distribution_type == "Random":
+            initial_grid_2d = selected_initial_state_function_raw(current_N, current_N, current_N)
+        else:
+            initial_grid_2d = initial_state_func_eval(initial_grid_2d_X, initial_grid_2d_Y)
+
+        sub_sv_init_flat = initial_grid_2d.flatten().astype(np.complex128)
+        full_initial_sv_host = np.zeros(N_sub_per_rank, dtype=np.complex128)
+        num_computational_states = current_N * current_N
+
+        if len(sub_sv_init_flat) == num_computational_states:
+            if num_computational_states <= N_sub_per_rank:
+                full_initial_sv_host[:num_computational_states] = sub_sv_init_flat
+            else:
+                print(f"Error: Grid data {num_computational_states} > N_sub_per_rank {N_sub_per_rank}")
+                return None, None
+        else:
+            print(f"Warning: Initial state size {len(sub_sv_init_flat)}!= expected {num_computational_states}")
+            fill_len = min(len(sub_sv_init_flat), num_computational_states, N_sub_per_rank)
+            full_initial_sv_host[:fill_len] = sub_sv_init_flat[:fill_len]
             
-    points_2d = np.vstack([X.ravel(), Y.ravel()]).T
-    tri = Delaunay(points_2d)
-    
-    fig = go.Figure()
-
-    for i in range(M):
-        time_value = (time_indices[i] / (Nt-1)) * t_max if Nt > 1 else 0
-        z_data = U_slider[i, :, :].flatten()
-        fig.add_trace(
-            go.Mesh3d(
-                x=X.flatten(), y=Y.flatten(), z=z_data,
-                i=tri.simplices[:, 0], j=tri.simplices[:, 1], k=tri.simplices[:, 2],
-                intensity=z_data,
-                colorscale='Viridis',
-                # --- FIX: Set the intensity range to the global min/max for a consistent color bar ---
-                cmin=z_min,
-                cmax=z_max,
-                # -------------------------------------------------------------------------------------
-                name=f'Time: {time_value:.2f}s',
-                showscale=True if i == 0 else False,
-                visible=(i == 0)
+        rank_slice_init = to_cupy_array(initial_state_val)
+        print(f'Rank {rank}: Initializing state with {distribution_type} (ux={ux_input}, uy={uy_input})...')
+        cp.cuda.runtime.memcpy(rank_slice_init.data.ptr, full_initial_sv_host.ctypes.data, full_initial_sv_host.nbytes, cp.cuda.runtime.memcpyHostToDevice)
+        print(f'Rank {rank}: Initial state copied. Size: {len(sub_sv_init_flat)}. N_sub_per_rank: {N_sub_per_rank}')
+
+        # Explicitly save initial state (t=0)
+        qlbm_obj.write_state(initial_state_val, "0")
+
+        print("Starting QLBM evolution...")
+        # Define specific timesteps to save for the slider
+        timesteps_for_slider = # T-1 is the last t_iter
+        qlbm_obj.run_evolution(initial_state_val, T, timesteps_to_save=timesteps_for_slider)
+        print("QLBM evolution complete.")
+
+        print("Generating plots with Plotly...")
+        downsampled_N = current_N // downsampling_factor
+        if downsampled_N == 0 and current_N > 0:
+            downsampled_N = 1
+        elif current_N == 0:
+            print("Error: current_N is zero before Plotly stage.")
+            return None, None
+
+        # Load data for specific time steps for interactive plot
+        # These correspond to the filenames saved: 0, T//4, 3*T//4, T-1
+        time_steps_to_load =
+        data_frames =
+        actual_timesteps_loaded =
+        for t in time_steps_to_load:
+            file_path = intermediate_folder_path / f"{t}_{rank}.npy"
+            if file_path.exists():
+                sol_loaded = np.load(file_path)
+                if sol_loaded.size == downsampled_N * downsampled_N:
+                    Z_data = np.reshape(sol_loaded, (downsampled_N, downsampled_N))
+                    data_frames.append(Z_data)
+                    actual_timesteps_loaded.append(t)
+                else:
+                    print(f"Warning: File {file_path} size {sol_loaded.size}!= expected {downsampled_N*downsampled_N}. Skipping.")
+            else:
+                print(f"Warning: File {file_path} not found. Skipping.")
+
+        if not data_frames:
+            print("Error: No data frames loaded for interactive plot.")
+            return None, None
+
+        x_coords_plot = np.linspace(-10, 10, downsampled_N)
+        y_coords_plot = np.linspace(-10, 10, downsampled_N)
+
+        # Calculate global min/max for consistent scaling
+        z_min = min([np.min(Z) for Z in data_frames])
+        z_max = max([np.max(Z) for Z in data_frames])
+        if z_max == z_min:
+            z_max += 1e-9
+
+        # Create interactive Plotly figure with slider
+        fig = go.Figure()
+
+        for i, Z in enumerate(data_frames):
+            fig.add_trace(
+                go.Surface(
+                    z=Z, x=x_coords_plot, y=y_coords_plot,
+                    colorscale='Viridis',
+                    cmin=z_min, cmax=z_max,
+                    name=f'Time: {actual_timesteps_loaded[i]}',
+                    showscale=(i == 0) # Show color scale only for the first trace
+                )
             )
-        )
 
-    # Slider creation logic is unchanged
-    steps = []
-    for i in range(len(fig.data)):
-        time_value = (time_indices[i] / (Nt-1)) * t_max if Nt > 1 else 0
-        step = dict(
-            method="update",
-            args=[{"visible": [False] * len(fig.data)}],
-            label=f"{time_value:.2f}s"
+        steps =
+        for i in range(len(data_frames)):
+            step = dict(
+                method="update",
+                args=[{"visible": [False] * len(data_frames)}],
+                label=f"Time: {actual_timesteps_loaded[i]}"
+            )
+            step["args"]["visible"][i] = True
+            steps.append(step)
+
+        sliders =
+
+        fig.update_layout(
+            title='QLBM Simulation - Density Evolution',
+            scene=dict(
+                xaxis_title='X',
+                yaxis_title='Y',
+                zaxis_title='Density',
+                xaxis=dict(range=[x_coords_plot, x_coords_plot[-1]]),
+                yaxis=dict(range=[y_coords_plot, y_coords_plot[-1]]),
+                zaxis=dict(range=[z_min, z_max]),
+            ),
+            sliders=sliders,
+            width=1000, # Increased width
+            height=900 # Increased height
         )
-        step["args"][0]["visible"][i] = True
-        steps.append(step)
-
-    sliders = [dict(active=0, currentvalue={"prefix": "Time: "}, pad={"t": 50}, steps=steps)]
-
-    # --- FIX: Update the layout to use the fixed axis ranges ---
-    fig.update_layout(
-        title=f'2D Heat Eq — init={initial}, bc={bc}',
-        scene=dict(
-            xaxis_title='X',
-            yaxis_title='Y',
-            zaxis_title='Temperature',
-            # Set the axis ranges to be fixed across all animation frames
-            xaxis=dict(range=[0, Lx]),
-            yaxis=dict(range=[0, Ly]),
-            zaxis=dict(range=[z_min, z_max]),
-        ),
-        sliders=sliders
+
+        # GIF generation logic removed as per request
+        #... (removed all GIF related code)...
+
+        return fig # Return only the interactive Plotly figure
+
+# Gradio Interface Definition
+def qlbm_gradio_interface(num_reg_qubits_input: int, timescale_input: int, distribution_type_param: str, ux_param: float, uy_param: float, velocity_field_type_param: str):
+    num_reg_qubits_val = int(num_reg_qubits_input)
+    timescale_val = int(timescale_input)
+    ux_val = float(ux_param)
+    uy_val = float(uy_param)
+
+    print(f"Gradio Interface: num_reg_qubits={num_reg_qubits_val}, T={timescale_val}, Distribution={distribution_type_param}, ux={ux_val}, uy={uy_val}, VelocityFieldType={velocity_field_type_param}")
+
+    plot_fig = simulate_qlbm_and_animate( # Only expecting plot_fig now
+        num_reg_qubits=num_reg_qubits_val,
+        T=timescale_val,
+        distribution_type=distribution_type_param,
+        ux_input=ux_val,
+        uy_input=uy_val,
+        velocity_field_type=velocity_field_type_param # Pass the new dummy parameter
     )
-    # -------------------------------------------------------------
-    
-    return fig
 
+    if plot_fig is None:
+        gr.Warning("Simulation or plotting failed. Please check console for errors.")
+        return None
+    return plot_fig # Return only the interactive Plotly figure
 
-# --- Gradio Interface Function with Error Handling (No changes here) ---
-def gradio_interface(lx, ly, t_max, m_steps, gamma, nx, ny, initial, bc):
-    try:
-        nx, ny, m_steps = int(nx), int(ny), int(m_steps)
-        fig = solve_and_plot_interactive(
-            Lx=lx, Ly=ly, t_max=t_max, M=m_steps, Gamma=gamma, Nx=nx, Ny=ny,
-            initial=initial, bc=bc
-        )
-        return fig
-    except Exception as e:
-        error_text = traceback.format_exc()
-        print(error_text)
-        error_fig = go.Figure().update_layout(
-            title_text="⚠️ Application Error",
-            annotations=[dict(text=f"An error occurred: {e}", showarrow=False)]
-        )
-        return error_fig
+with gr.Blocks(theme=gr.themes.Soft(), title="QLBM Simulation with Plotly") as qlbm_demo:
+    gr.Markdown(
+    """
+    # ⚛️ Quantum Lattice Boltzmann Method (QLBM) Simulator (Plotly Animation)
+    Welcome to the Quantum Lattice Boltzmann Method (QLBM) simulator! This version uses Plotly for 3D animation and interactive plots.
 
-# --- Gradio UI Definition (No changes here) ---
-with gr.Blocks(theme=gr.themes.Soft(), title="2D Heat Simulator") as demo:
-    gr.Markdown("# ♨️ 2D Heat Equation Simulator")
-    gr.Markdown("Adjust parameters and click 'Run' to generate an interactive plot directly in your browser.")
+    **How this Simulator Works:**
+    This simulator implements a D2Q5 model on a quantum computer simulator (CUDA-Q).
+    - Control grid size (via Number of Register Qubits: $N=2^{\text{num_reg_qubits}}$).
+    - Set total simulation time (Timescale T).
+    - Choose initial distribution.
+    - Set advection velocities `ux` and `uy`.
+    The simulation generates an interactive Plotly figure with a slider for selected time steps.
 
+    **Note:** Higher qubit counts and longer timescales are computationally intensive. Advection velocities should be small (e.g., < 0.3).
+    The Plotly figure allows interactive exploration of specific time steps.
+    """
+    )
     with gr.Row():
         with gr.Column(scale=1):
             gr.Markdown("## Simulation Parameters")
-            lx_slider = gr.Slider(0.1, 5.0, 1.0, 0.1, label="Lx")
-            ly_slider = gr.Slider(0.1, 5.0, 1.0, 0.1, label="Ly")
-            nx_slider = gr.Slider(10, 80, 40, 1, label="Nx (Grid Points X)")
-            ny_slider = gr.Slider(10, 80, 40, 1, label="Ny (Grid Points Y)")
-            t_slider = gr.Slider(0.01, 2.0, 0.5, 0.01, label="t_max")
-            gamma_slider = gr.Slider(0.001, 1.0, 0.1, 0.001, label="Gamma")
-            m_slider = gr.Slider(10, 80, 30, 1, label="M (Time Steps)")
-            
-            with gr.Row():
-                initial_dropdown = gr.Dropdown(["gaussian", "random", "sinusoidal", "step"], value="gaussian", label="Initial")
-                bc_dropdown = gr.Dropdown(["dirichlet", "neumann", "periodic"], value="dirichlet", label="Boundary")
+            num_reg_qubits_slider = gr.Slider(
+                minimum=2, maximum=10, value=8, step=1,
+                label="Number of Register Qubits (num_reg_qubits)",
+                info="Grid N = 2^num_reg_qubits. Max 10 (Note: >8 slow; >9 may hit simulator/memory limits on free tiers)."
+            )
+            timescale_slider = gr.Slider(
+                minimum=0, maximum=2000, value=100, step=10,
+                label="Timescale (T)", info="Total number of timesteps. Max 2000."
+            )
+
+            # Group 1: Initial Conditions
+            with gr.Accordion("Initial Conditions", open=True):
+                distribution_options =
+                distribution_type_input = gr.Radio(
+                    choices=distribution_options, value="Sine Wave (Original)",
+                    label="Initial Distribution Type", info="Select the initial pattern of the substance."
+                )
+
+            # Group 2: Velocity Fields
+            with gr.Accordion("Velocity Fields", open=True):
+                velocity_field_options = # Dummy options
+                velocity_field_type_input = gr.Radio(
+                    choices=velocity_field_options, value="Uniform",
+                    label="Velocity Field Type", info="Select the type of background velocity field."
+                )
+                ux_slider = gr.Slider(
+                    minimum=-0.4, maximum=0.4, value=0.2, step=0.01,
+                    label="Advection Velocity ux", info="x-component of background advection."
+                )
+                uy_slider = gr.Slider(
+                    minimum=-0.4, maximum=0.4, value=0.15, step=0.01,
+                    label="Advection Velocity uy", info="y-component of background advection."
+                )
             
-            run_btn = gr.Button("Run Simulation", variant="primary")
+            run_qlbm_btn = gr.Button("Run QLBM Simulation", variant="primary")
 
-        with gr.Column(scale=3):
-            plot_output = gr.Plot(label="Interactive Heatmap")
+        with gr.Column(scale=2):
+            # Removed gr.Image for GIF
+            # qlbm_plot_output = gr.Image(label="QLBM Simulation Animation (GIF)", type="filepath", height=900)
+            qlbm_interactive_plot = gr.Plot(label="Interactive Density Plot with Slider")
 
-    inputs_list = [lx_slider, ly_slider, t_slider, m_slider, gamma_slider,
-                   nx_slider, ny_slider, initial_dropdown, bc_dropdown]
-    
-    run_btn.click(fn=gradio_interface, inputs=inputs_list, outputs=plot_output)
-    
+    qlbm_inputs_list = [num_reg_qubits_slider, timescale_slider, distribution_type_input, ux_slider, uy_slider, velocity_field_type_input]
+    run_qlbm_btn.click(
+        fn=qlbm_gradio_interface,
+        inputs=qlbm_inputs_list,
+        outputs=[qlbm_interactive_plot] # Only interactive plot
+    )
     gr.Examples(
-        examples=[
-            [1.0, 1.0, 0.5, 30, 0.1, 30, 30, "gaussian", "dirichlet"],
-            [2.0, 1.0, 1.0, 40, 0.05, 40, 20, "sinusoidal", "periodic"],
-        ],
-        inputs=inputs_list,
-        outputs=plot_output,
-        fn=gradio_interface,
+        examples=,
+            [6, 50, "Gaussian", 0.1, 0.05, "Uniform"],
+           ,
+        inputs=qlbm_inputs_list,
+        outputs=[qlbm_interactive_plot], # Only interactive plot
+        fn=qlbm_gradio_interface,
         cache_examples=False
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    try:
+        cudaq.set_target('nvidia', option='fp64')
+        print(f"CUDA-Q Target successfully set to: {cudaq.get_target().name}")
+    except Exception as e_target:
+        print(f"Warning: Could not set CUDA-Q target to 'nvidia'. Error: {e_target}")
+        print(f"Current CUDA-Q Target: {cudaq.get_target().name}. Performance may be affected.")
+    
+    qlbm_demo.launch()