fluid.py

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 plotly.io as pio
from sympy import sympify, symbols, lambdify
from gradio_litmodel3d import LitModel3D
import zipfile

# Set Plotly engine for image export
try:
    pio.kaleido.scope.mathjax = None
except AttributeError:
    pass

# Existing functions (bin_to_gray, gray_to_bin, etc.) remain unchanged
def bin_to_gray(bin_s):
    XOR=lambda x,y: (x or y) and not (x and y)
    gray_s=bin_s[0]
    for i in range(len(bin_s)-1):
        c_bool=XOR(bool(int(bin_s[i])),bool(int(bin_s[i+1])))
        gray_s+=str(int(c_bool))
    return gray_s

def gray_to_bin(gray_s):
  XOR=lambda x,y: (x or y) and not (x and y)
  bin_s=gray_s[0]
  for i in range(len(gray_s)-1):
    c_bool=XOR(bool(int(bin_s[i])),bool(int(gray_s[i+1])))
    bin_s+=str(int(c_bool))
  return bin_s

def bin_to_int(bin_s):
  return int(bin_s,2)

def int_to_bin(i,pad):
  return bin(i)[2:].zfill(pad)

def fwht_approx(f,N,num_points_per_dim,threshold=1e-10):
  linear_block_size=int(N//num_points_per_dim)
  num_angles_per_block=int(np.log2(linear_block_size))

  thetas={}

  for k in range(num_points_per_dim):
    for j in range(num_points_per_dim):
      for i in range(num_points_per_dim):

        avg_f=2*np.arccos(f(i*linear_block_size+(linear_block_size-1)/2,j*linear_block_size+(linear_block_size-1)/2,k*linear_block_size+(linear_block_size-1)/2))
        thetas[k*(N**2)*linear_block_size+j*N*linear_block_size+i*linear_block_size]=avg_f

        slope_x=(2*np.arccos(f(i*linear_block_size,j*linear_block_size+(linear_block_size-1)/2,k*linear_block_size+(linear_block_size-1)/2))-2*np.arccos(f(((i+1)%N)*linear_block_size,j*linear_block_size+(linear_block_size-1)/2,k*linear_block_size+(linear_block_size-1)/2)))/linear_block_size
        slope_y=(2*np.arccos(f(i*linear_block_size+(linear_block_size-1)/2,j*linear_block_size,k*linear_block_size+(linear_block_size-1)/2))-2*np.arccos(f(i*linear_block_size+(linear_block_size-1)/2,((j+1)%N)*linear_block_size,k*linear_block_size+(linear_block_size-1)/2)))/linear_block_size
        slope_z=(2*np.arccos(f(i*linear_block_size+(linear_block_size-1)/2,j*linear_block_size+(linear_block_size-1)/2,k*linear_block_size))-2*np.arccos(f(i*linear_block_size+(linear_block_size-1)/2,j*linear_block_size+(linear_block_size-1)/2,((k+1)%N)*linear_block_size)))/linear_block_size

        for m in range(num_angles_per_block):
          thetas[k*(N**2)*linear_block_size+j*N*linear_block_size+i*linear_block_size + 2**m]=slope_x*(2**(m-1))
          thetas[k*(N**2)*linear_block_size+j*N*linear_block_size+i*linear_block_size + N*(2**m)]=slope_y*(2**(m-1))
          thetas[k*(N**2)*linear_block_size+j*N*linear_block_size+i*linear_block_size + (N**2)*(2**m)]=slope_z*(2**(m-1))

  h = linear_block_size
  while h < N**3:
    for i in range(0, N**3, h * 2):
      if (i//N)%linear_block_size!=0:
        continue
      if (i//(N**2))%linear_block_size!=0:
        continue
      j=i
      while j<i+h:
        index=j
        x = thetas[index]
        y = thetas[index + h]
        thetas[index] = (x + y)/2
        thetas[index + h] = (x - y)/2

        for ax in range(3):
          for m in range(num_angles_per_block):
            index = j + (N**ax) * (2**m)
            x = thetas[index]
            y = thetas[index + h]
            thetas[index] = (x + y)/2
            thetas[index + h] = (x - y)/2

        j+=linear_block_size
        if (j//N)%linear_block_size==1:
          j+=(linear_block_size-1)*N
        if (j//(N**2))%linear_block_size==1:
          j+=(linear_block_size-1)*(N**2)

    h *= 2
    if h==N:
      h=N*linear_block_size
    if h==N**2:
      h=(N**2)*linear_block_size

  return [theta for theta in thetas.values() if abs(theta)>threshold],[key for key in thetas.keys() if abs(thetas[key])>threshold]

def get_circuit_inputs(f,num_reg_qubits,num_points_per_dim):
  theta_vec,indices=fwht_approx(f,2**num_reg_qubits,num_points_per_dim)
  circ_pos=[]
  for ind in indices:
    circ_pos+=[bin_to_int(gray_to_bin(int_to_bin(ind,num_reg_qubits*3)))]

  sorted_theta_vec=sorted(zip(theta_vec,circ_pos),key=lambda el:el[1])
  ctrls=[]

  current_bs="0"*(3*num_reg_qubits)
  for el in sorted_theta_vec:
    new_bs=bin_to_gray(int_to_bin((el[1])%(2**(3*num_reg_qubits)),(3*num_reg_qubits)))
    ctrls += [[i for i, (char1, char2) in enumerate(zip(current_bs, new_bs)) if char1 != char2]]
    current_bs=new_bs
  new_bs="0"*(3*num_reg_qubits)
  ctrls += [[i for i, (char1, char2) in enumerate(zip(current_bs, new_bs)) if char1 != char2]]

  ctrls_flat_list=[]
  for ctrl_list in ctrls:
    ctrls_flat_list+=[len(ctrl_list)]+ctrl_list

  return [el[0] for el in sorted_theta_vec]+[0.0],ctrls_flat_list

# Simulation functions remain unchanged
def simulate_qlbm_and_animate(num_reg_qubits: int, T: int, distribution_type: str, velocity_field: str, vx_input: float, vy_input: float, boundary_condition: str):
    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)
    NUM_ANIMATION_FRAMES = 40

    if T == 0:
        time_steps = [0]
    else:
        num_points = min(T + 1, NUM_ANIMATION_FRAMES)
        time_steps = np.linspace(start=0, stop=T, num=num_points, dtype=int)
        time_steps = sorted(list(set(time_steps)))

    if distribution_type == "Sinusoidal":
        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
    else:
        print(f"Warning: Unknown distribution type '{distribution_type}'. Defaulting to Sinusoidal.")
        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)

    if velocity_field == "Uniform":
        pass

    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, vx=0.2, vy=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.vx = vx
                self.vy = vy
                self.u = np.array([0, self.vx, self.vx, self.vy, self.vy])
                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 = 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[0])
                        elif i == 1:
                            x.ctrl(q[0], q[1])
                        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[2], qdir[1], qdir[0])
                    d2q5_tstep(q, nqx, nqy, nq_dir_val, dirs_i_val)
                    prep_op(qdir[2], qdir[1], qdir[0])

                def run_timestep_func(vec_arg, hadamard=False):
                    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, time_steps_to_save, observable=False):
                current_state_val = initial_state_arg
                save_times = set(time_steps_to_save)
                if 0 in save_times:
                    self.write_state(current_state_val, '0')
                for t_iter in range(total_timesteps):
                    if (t_iter + 1) in save_times:
                        next_state_val = self.run_timestep(current_state_val)
                        self.write_state(next_state_val, str(t_iter + 1))
                        current_state_val = next_state_val
                    else:
                        current_state_val = self.run_timestep(current_state_val)
                    cp.get_default_memory_pool().free_all_blocks()
                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

        if boundary_condition == "Periodic":
            pass

        downsampling_factor = 2**5
        if current_N == 0:
            print("Error: current_N is zero. num_reg_qubits likely too small.")
            return None, None # Modified return
        if current_N < downsampling_factor:
            downsampling_factor = current_N if current_N > 0 else 1

        qlbm_obj = QLBMAdvecDiffD2Q5_new(vx=vx_input, vy=vy_input)
        initial_state_val = cudaq.get_state(alloc_kernel, num_qubits_total)

        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)

        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)
        norm = np.linalg.norm(sub_sv_init_flat)
        if norm > 0:
            sub_sv_init_flat /= norm
        else:
            print("Error: Initial state norm is zero.")
            return None, None # Modified return
        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 # Modified return
        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]

        rank_slice_init = to_cupy_array(initial_state_val)
        print(f'Rank {rank}: Initializing state with {distribution_type} (vx={vx_input}, vy={vy_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}')

        print("Starting QLBM evolution...")
        qlbm_obj.run_evolution(initial_state_val, T, time_steps)
        print("QLBM evolution complete.")

        print("Generating interactive plot 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 # Modified return

        data_frames = []
        actual_timesteps = []
        for t in time_steps:
            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.append(t)
                    print(f"Time {t}: Min={np.min(Z_data)}, Max={np.max(Z_data)}, Mean={np.mean(Z_data)}")
                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 plotting.")
            return None, None # Modified return

        x_coords_plot = np.linspace(0, 1, downsampled_N)
        y_coords_plot = np.linspace(0, 1, downsampled_N)

        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

        fig = go.Figure()
        
        # Store individual frames for download
        plotly_json_frames = []

        for i, Z in enumerate(data_frames):
            frame_trace = go.Surface(
                z=Z, x=x_coords_plot, y=y_coords_plot,
                colorscale='Blues',
                cmin=z_min, cmax=z_max,
                name=f'Time: {actual_timesteps[i]}',
                showscale=False
            )
            fig.add_trace(frame_trace)
            
            # Create a figure for the individual frame and convert to JSON
            single_frame_fig = go.Figure(data=[frame_trace], layout=fig.layout)
            single_frame_fig.update_layout(
                title=f"Time: {actual_timesteps[i]}",
                scene=dict(
                    xaxis_title='X',
                    yaxis_title='Y',
                    zaxis_title='Density',
                    xaxis=dict(range=[x_coords_plot[0], x_coords_plot[-1]]),
                    yaxis=dict(range=[y_coords_plot[0], y_coords_plot[-1]]),
                    zaxis=dict(range=[z_min, z_max]),
                )
            )
            plotly_json_frames.append(single_frame_fig.to_json())


        for trace in fig.data[1:]:
            trace.visible = False

        steps = []
        for i in range(len(data_frames)):
            step = dict(
                method="update",
                args=[{"visible": [False] * len(data_frames)}],
                label=f"Time: {actual_timesteps[i]}"
            )
            step["args"][0]["visible"][i] = True
            steps.append(step)

        sliders = [dict(active=0, currentvalue={"prefix": "Time: "}, pad={"t": 50}, steps=steps)]

        # fig.update_layout(
        #     title='', # Removed graph title
        #     scene=dict(
        #         xaxis_title='X',
        #         yaxis_title='Y',
        #         zaxis_title='Density',
        #         xaxis=dict(range=[x_coords_plot[0], x_coords_plot[-1]]),
        #         yaxis=dict(range=[y_coords_plot[0], y_coords_plot[-1]]),
        #         zaxis=dict(range=[z_min, z_max]),
        #     ),
        #     sliders=sliders,
        #     width=800,
        #     height=700
        # )
        fig.update_layout(
            title='', # Removed graph title
            scene=dict(
                xaxis_title='X',
                yaxis_title='Y',
                zaxis_title='Density',
                xaxis=dict(visible=False),
                yaxis=dict(visible=False),
                zaxis=dict(visible=False),
            ),
            sliders=sliders,
            width=800,
            height=700
        )
        return fig, plotly_json_frames # Modified return