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 jthreshold],[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 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[0], x.shape[1]) * 1.5 + 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 == "User": pass elif velocity_field == "Shear": vx_input = vx_input * (current_N / 2) elif velocity_field == "TGV": vx_input = vx_input * np.sin(np.pi * current_N / 10) vy_input = vy_input * np.cos(np.pi * current_N / 10) elif velocity_field == "Swirl": vx_input = vx_input * np.cos(np.pi * current_N / 5) vy_input = vy_input * np.sin(np.pi * current_N / 5) 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 elif boundary_condition == "Dirichlet": pass elif boundary_condition == "Neumann": 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) 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) 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 def simulate_qlbm_3D_and_animate(num_reg_qubits: int, T: int, distribution_type: str, vx_input, vy_input, vz_input, boundary_condition: str): num_anc = 3 num_qubits_total = 3 * 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) # Simplified time steps for 3D since slider steps are removed NUM_ANIMATION_FRAMES_3D = 10 # Default number of frames if no specific slider steps if T == 0: time_steps = [0] else: num_points = min(T + 1, NUM_ANIMATION_FRAMES_3D) 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, z, N_val_func: \ np.sin(x * 2 * np.pi / N_val_func) * \ np.sin(y * 2 * np.pi / N_val_func) * \ np.sin(z * 2 * np.pi / N_val_func) + 1 elif distribution_type == "Gaussian": selected_initial_state_function_raw = lambda x, y, z, 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) + (z - 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, z, N_val_func: \ np.sin(x * 2 * np.pi / N_val_func) * \ np.sin(y * 2 * np.pi / N_val_func) * \ np.sin(z * 2 * np.pi / N_val_func) + 1 initial_state_func_eval = lambda i:\ selected_initial_state_function_raw(i%current_N,(i//current_N)%current_N,i//(current_N**2),current_N)*(i<(current_N**3)).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 QLBMAdvecDiffD3Q7_new: def __init__(self,vx,vy,vz) -> None: self.dim = 3 self.ndir = 7 self.nq_dir = math.ceil(np.log2(self.ndir)) self.dirs=[] for dir_int in range(self.ndir): if dir_int==4: dir_bin="111" else: dir_bin = f"{dir_int:b}".zfill(self.nq_dir) self.dirs.append(dir_bin) self.cs = 1/np.sqrt(3) self.ux = lambda x,y,z: vx(x,y,z)/self.cs**2 self.uy = lambda x,y,z: vy(x,y,z)/self.cs**2 self.uz = lambda x,y,z: vz(x,y,z)/self.cs**2 self.create_circuit() def create_circuit(self): print("Creating circuit") x_coeffs,x_coeff_var_indices=get_circuit_inputs(lambda x,y,z: ((1+self.ux(x/current_N,y/current_N,z/current_N))/2)**0.5,num_reg_qubits,min(current_N,32)) y_coeffs,y_coeff_var_indices=get_circuit_inputs(lambda x,y,z: ((1+self.uy(x/current_N,y/current_N,z/current_N))/2)**0.5,num_reg_qubits,min(current_N,32)) z_coeffs,z_coeff_var_indices=get_circuit_inputs(lambda x,y,z: ((1+self.uz(x/current_N,y/current_N,z/current_N))/2)**0.5,num_reg_qubits,min(current_N,32)) x_coeffs_,x_coeff_var_indices_=get_circuit_inputs(lambda x,y,z: 0 if (1+self.ux((x-1)/current_N,y/current_N,z/current_N))==0 else \ ((1+self.ux((x-1)/current_N,y/current_N,z/current_N))/(2+self.ux((x-1)/current_N,y/current_N,z/current_N)-self.ux((x+1)/current_N,y/current_N,z/current_N)))**0.5,num_reg_qubits,min(current_N,32)) y_coeffs_,y_coeff_var_indices_=get_circuit_inputs(lambda x,y,z: 0 if (1+self.uy(x/current_N,(y-1)/current_N,z/current_N))==0 else \ ((1+self.uy(x/current_N,(y-1)/current_N,z/current_N))/(2+self.uy(x/current_N,(y-1)/current_N,z/current_N)-self.uy(x/current_N,(y+1)/current_N,z/current_N)))**0.5,num_reg_qubits,min(current_N,32)) z_coeffs_,z_coeff_var_indices_=get_circuit_inputs(lambda x,y,z: 0 if (1+self.uz(x/current_N,y/current_N,(z-1)/current_N))==0 else \ ((1+self.uz(x/current_N,y/current_N,(z-1)/current_N))/(2+self.uz(x/current_N,y/current_N,(z-1)/current_N)-self.uz(x/current_N,y/current_N,(z+1)/current_N)))**0.5,num_reg_qubits,min(current_N,32)) unprep1_coeffs,unprep1_coeff_var_indices=get_circuit_inputs(lambda x,y,z:\ (1/3**0.5)*(1+(self.ux((x-1)/current_N,y/current_N,z/current_N)-self.ux((x+1)/current_N,y/current_N,z/current_N))/2)**0.5,num_reg_qubits,min(current_N,32)) unprep2_coeffs, unprep2_coeff_var_indices = get_circuit_inputs(lambda x, y, z: ((1 + (self.uy(x/current_N, (y-1)/current_N, z/current_N) - self.uy(x/current_N, (y+1)/current_N, z/current_N))/2) /(2 - (self.ux((x-1)/current_N, y/current_N, z/current_N) - self.ux((x+1)/current_N, y/current_N, z/current_N))/2))**0.5, num_reg_qubits, min(current_N, 32)) print("Generated angles") v=np.pad([1/4, 1/4, 0, 1/4, 0, 1/4, 0],(0,2**num_anc - self.ndir)) v=v**0.5 v[0]+=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): dirs_i_list=[] for dir_ in dirs: dirs_i=[(int(c)) for c in dir_] dirs_i_list+=dirs_i[::-1] return dirs_i_list self.dirs_i_list=collisionOp(self.dirs) print("Generated dirs_i_list") @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, nqz: int, nq_dir: int, dirs_i: list[int]): qx=q[0:nqx] qy=q[nqx:nqx+nqy] qz=q[nqx+nqy:nqx+nqy+nqz] qdir=q[nqx+nqy+nqz:nqx+nqy+nqz+nq_dir] i=2 b_list=dirs_i[i*nq_dir:(i+1)*nq_dir] for j in range(nq_dir): b=b_list[j] if b==0: x(qdir[j]) cudaq.control(lshift,qdir,qx,nqx) for j in range(nq_dir): b=b_list[j] if b==0: x(qdir[j]) i=1 b_list=dirs_i[i*nq_dir:(i+1)*nq_dir] for j in range(nq_dir): b=b_list[j] if b==0: x(qdir[j]) cudaq.control(rshift,qdir,qx,nqx) for j in range(nq_dir): b=b_list[j] if b==0: x(qdir[j]) i=4 b_list=dirs_i[i*nq_dir:(i+1)*nq_dir] for j in range(nq_dir): b=b_list[j] if b==0: x(qdir[j]) cudaq.control(lshift,qdir,qy,nqy) for j in range(nq_dir): b=b_list[j] if b==0: x(qdir[j]) i=3 b_list=dirs_i[i*nq_dir:(i+1)*nq_dir] for j in range(nq_dir): b=b_list[j] if b==0: x(qdir[j]) cudaq.control(rshift,qdir,qy,nqy) for j in range(nq_dir): b=b_list[j] if b==0: x(qdir[j]) i=6 b_list=dirs_i[i*nq_dir:(i+1)*nq_dir] for j in range(nq_dir): b=b_list[j] if b==0: x(qdir[j]) cudaq.control(lshift,qdir,qz,nqz) for j in range(nq_dir): b=b_list[j] if b==0: x(qdir[j]) i=5 b_list=dirs_i[i*nq_dir:(i+1)*nq_dir] for j in range(nq_dir): b=b_list[j] if b==0: x(qdir[j]) cudaq.control(rshift,qdir,qz,nqz) for j in range(nq_dir): b=b_list[j] if b==0: x(qdir[j]) @cudaq.kernel def d2q5_tstep_wrapper(state: cudaq.State,nqx:int,nqy:int,nqz:int,nq_dir:int,dirs_i:list[int],\ x_coeff_var_indices:list[int],x_coeffs:list[float],\ y_coeff_var_indices:list[int],y_coeffs:list[float],\ z_coeff_var_indices:list[int],z_coeffs:list[float],\ x_coeff_var_indices_:list[int],x_coeffs_:list[float],\ y_coeff_var_indices_:list[int],y_coeffs_:list[float],\ z_coeff_var_indices_:list[int],z_coeffs_:list[float],\ unprep1_coeff_var_indices:list[int],unprep1_coeffs:list[float],\ unprep2_coeff_var_indices:list[int],unprep2_coeffs:list[float]): q=cudaq.qvector(state) qdir=q[nqx+nqy+nqz:nqx+nqy+nqz+nq_dir] prep_op(qdir[2],qdir[1],qdir[0]) x.ctrl(qdir[0],qdir[1]) ind=0 coeff_ind=0 x(qdir[2]) while ind= 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 print("Circuit created") 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 * current_N: arr_flat = relevant_part_cupy[:current_N * current_N * current_N] if downsampling_factor > 1 and current_N > 0: arr_reshaped = arr_flat.reshape((current_N, current_N, current_N)) arr_downsampled = arr_reshaped[::downsampling_factor, ::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) print("Write state defined") 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: print("Writing first state") self.write_state(current_state_val, '0') for t_iter in range(total_timesteps): print("Running timestep") next_state_val = self.run_timestep(current_state_val) if (t_iter + 1) in save_times: print("Writing next state") self.write_state(next_state_val, str(t_iter + 1)) 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 if boundary_condition == "Periodic": pass elif boundary_condition == "Dirichlet": pass elif boundary_condition == "Neumann": pass downsampling_factor = 1 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 = QLBMAdvecDiffD3Q7_new(vx=vx_input, vy=vy_input, vz=vz_input) initial_state_val = cudaq.get_state(alloc_kernel, num_qubits_total) sub_sv_init_flat = initial_state_func_eval(np.arange(N_sub_per_rank)).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 ** 3 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* downsampled_N: data = np.reshape(sol_loaded, (downsampled_N, downsampled_N, downsampled_N)) data_frames.append(data) actual_timesteps.append(t) print(f"Time {t}: Min={np.min(data)}, Max={np.max(data)}, Mean={np.mean(data)}") else: print(f"Warning: File {file_path} size {sol_loaded.size} != expected {downsampled_N*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_coords_plot = np.linspace(0, 1, downsampled_N) Z_grid_mesh, Y_grid_mesh, X_grid_mesh = np.meshgrid(x_coords_plot, y_coords_plot, z_coords_plot, indexing='ij') data_min = min([np.min(data) for data in data_frames]) data_max = max([np.max(data) for data in data_frames]) if data_max == data_min: data_max += 1e-9 fig = go.Figure() # Store individual frames for download plotly_json_frames = [] for i, output_data in enumerate(data_frames): frame_trace = go.Isosurface( x=X_grid_mesh.flatten(), y=Y_grid_mesh.flatten(), z=Z_grid_mesh.flatten(), value=output_data.flatten(), isomin=data_min, isomax=data_max, opacity=0.4, # needs to be small to see through all surfaces surface_count=7, # needs to be a large number for good volume rendering, caps=dict(x_show=False, y_show=False, z_show=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='Z', 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_coords_plot[0], z_coords_plot[-1]]), ) ) 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='Z', 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_coords_plot[0], z_coords_plot[-1]]), ), sliders=sliders, width=800, height=700 ) return fig, plotly_json_frames # Modified return