import yaml import os import numpy as np from scipy.optimize import root import csv import matplotlib.pyplot as plt def load_yaml(filepath): """Loads a YAML file and returns the dictionary.""" if not os.path.exists(filepath): raise FileNotFoundError(f"Missing {filepath}! Please check your paths.") with open(filepath, 'r') as file: return yaml.safe_load(file) def calculate_brine_density(mass_ions_g, mass_total_g): """ Calculates brine density (kg/L) based on the mass fraction of salt. """ if mass_total_g == 0: return 1.0 mass_fraction = mass_ions_g / mass_total_g # Simple empirical approximation: water density + salt mass fraction return 1.0 + mass_fraction def initialize_feed_mass(ions_dict, feed_volume_L): """Converts feed concentrations (mg/L) into absolute masses (grams).""" masses_g = {} total_tds_mg_L = sum(details.get('value', 0.0) for details in ions_dict.values()) for ion, details in ions_dict.items(): conc_mg_L = details.get('value', 0.0) masses_g[ion] = (conc_mg_L * feed_volume_L) / 1000.0 # We estimate initial mass using our density function initial_mass_fraction = total_tds_mg_L / 1_000_000.0 density_kg_L = 1.0 + initial_mass_fraction total_solution_mass_g = feed_volume_L * density_kg_L * 1000.0 total_ion_mass_g = sum(masses_g.values()) masses_g['H2O'] = total_solution_mass_g - total_ion_mass_g return masses_g def run_forward_mode(params, ions_dict, feed_volume, solvent_mass_g): """Runs the two-stage thermodynamic extraction simulation.""" # 1. Initialize Feed Mass feed_masses_g = initialize_feed_mass(ions_dict, feed_volume) # We need a fixed list of species to keep our arrays ordered species_list = ['H2O'] + sorted(ions_dict.keys()) # Create an initial guess: assume 0 grams of ions transfer, and 100g of water transfers initial_guesses = [100.0 if sp == 'H2O' else 0.0 for sp in species_list] # ========================================== # STAGE 1: EXTRACTION # ========================================== print("\n⚙️ Running Stage 1 Extraction (T_low)...") solution_1 = root( stage_1_residuals, x0=initial_guesses, args=(species_list, feed_masses_g, solvent_mass_g, params), method='hybr' ) if not solution_1.success: print("❌ Stage 1 Solver failed to converge:", solution_1.message) return print("✅ Stage 1 Equilibrium Reached!") # Unpack Stage 1 solution loaded_org_masses = {species: solution_1.x[i] for i, species in enumerate(species_list)} reject_aq_masses = {species: feed_masses_g[species] - solution_1.x[i] for i, species in enumerate(species_list)} print(f" -> Reject Brine left behind: {sum(reject_aq_masses.values()):.2f} g") print(f" -> Water/Salts pulled into solvent: {sum(loaded_org_masses.values()):.2f} g") # ========================================== # STAGE 2: RECOVERY # ========================================== print("\n⚙️ Running Stage 2 Recovery (T_high)...") initial_guesses_2 = [10.0 if sp == 'H2O' else 0.0 for sp in species_list] solution_2 = root( stage_2_residuals, x0=initial_guesses_2, args=(species_list, loaded_org_masses, solvent_mass_g, params), method='hybr' ) if not solution_2.success: print("❌ Stage 2 Solver failed to converge:", solution_2.message) return print("✅ Stage 2 Equilibrium Reached!") # Unpack Stage 2 solution final_org_masses = {species: solution_2.x[i] for i, species in enumerate(species_list)} yield_aq_masses = {species: loaded_org_masses[species] - solution_2.x[i] for i, species in enumerate(species_list)} # Convert masses back to volume and mg/L reject_vol_L, reject_conc = convert_mass_to_concentration(reject_aq_masses) yield_vol_L, yield_conc = convert_mass_to_concentration(yield_aq_masses) yield_tds = sum(yield_conc.values()) # Output Results print("\n==========================================") print("🏆 FINAL PROCESS OUTPUT") print("==========================================") print("\n--- Reject Brine ---") print(f" Volume: {reject_vol_L:.3f} L") print(f" Na: {reject_conc.get('Na', 0):.1f} mg/L") print(f" Cl: {reject_conc.get('Cl', 0):.1f} mg/L") print("\n--- Yield Water (Agricultural Product) ---") print(f" Volume: {yield_vol_L:.3f} L") print(f" Total TDS: {yield_tds:.1f} mg/L") for sp, conc in sorted(yield_conc.items()): print(f" {sp}: {conc:.2f} mg/L") print("\n--- Recycled Working Fluid (Dry) ---") print(f" Solvent (Base): {solvent_mass_g:.2f} g") print(f" Residual H2O: {final_org_masses['H2O']:.2f} g") def run_reverse_mode(target_yield_g, solvent_mass_g, s_low_max=0.833, steps=100): """ Sweeps S_low up to the 400% volumetric expansion limit (0.833 mass fraction) and calculates the required S_high to hit the target yield. Exports a .csv and saves a .png plot. """ print(f"\n⚙️ Running Reverse Mode Parameter Sweep...") print(f"Targeting {target_yield_g:.1f} g of water yield from {solvent_mass_g:.1f} g of solvent.") target_ratio = target_yield_g / solvent_mass_g # Calculate the absolute minimum S_low required to even hold the target yield # (If S_high was 0, what must S_low be?) s_low_min = target_ratio / (1.0 + target_ratio) if s_low_max <= s_low_min: print("❌ Error: The target yield requires more water capacity than your maximum limit allows.") return s_low_values = np.linspace(s_low_min + 0.001, s_low_max, steps) valid_s_low = [] s_high_values = [] for s_low in s_low_values: # Water held per gram of dry solvent at T_low water_ratio_low = s_low / (1.0 - s_low) # Water that MUST be held at T_high to expel the target amount water_ratio_high = water_ratio_low - target_ratio if water_ratio_high >= 0: s_high = water_ratio_high / (1.0 + water_ratio_high) valid_s_low.append(s_low) s_high_values.append(s_high) # 1. Export to CSV csv_filename = "tradeoff_data.csv" with open(csv_filename, mode='w', newline='') as file: writer = csv.writer(file) writer.writerow(["S_water_low_mass_fraction", "S_water_high_mass_fraction"]) for low, high in zip(valid_s_low, s_high_values): writer.writerow([f"{low:.4f}", f"{high:.4f}"]) print(f"✅ Data exported to {csv_filename}") # 2. Generate and Save Plot (Headless mode) plt.figure(figsize=(8, 6)) plt.plot(valid_s_low, s_high_values, 'b-', linewidth=2, label="Required Thermodynamic Path") # Shade the physically impossible region (below zero) plt.axhline(0, color='black', linewidth=1) plt.fill_between(valid_s_low, s_high_values, 0, alpha=0.1, color='blue') plt.title(f"Working Fluid Design Space\n(Target: {target_yield_g:.0f}g Yield per {solvent_mass_g:.0f}g Solvent)") plt.xlabel("Water Uptake at Extraction ($T_{low}$) [Mass Fraction]") plt.ylabel("Water Uptake at Recovery ($T_{high}$) [Mass Fraction]") plt.xlim([0.5, 0.85]) plt.ylim([0, max(s_high_values) * 1.1]) plt.grid(True, linestyle='--', alpha=0.7) plt.legend() png_filename = "tradeoff_plot.png" plt.savefig(png_filename, dpi=300, bbox_inches='tight') print(f"✅ Plot saved to {png_filename}") plt.close() # Prevents the plot from popping up and halting the script def run_empirical_mode(experiment_name, run_data): """ Parses lab data from a physical experiment, runs integrity checks, and calculates empirical water solubilities and volumetric salt partition coefficients. """ print(f"\n--------------------------------------------------") print(f"🔬 Analyzing: {experiment_name}") print(f"--------------------------------------------------") inputs = run_data['inputs'] outputs = run_data['measured_outputs'] if outputs['hydrated_wf_g'] is None: print(f"❌ Incomplete data for extraction stage. Cannot calculate thermodynamics.") return # ========================================== # 1. ABSOLUTE WATER UPTAKE # ========================================== dry_wf_in = inputs['wf2_base_g'] + inputs['acid_g'] water_extracted_g = outputs['hydrated_wf_g'] - dry_wf_in s_low = water_extracted_g / outputs['hydrated_wf_g'] water_trapped_g = outputs['dried_wf_g'] - dry_wf_in s_high = max(0.0, water_trapped_g / outputs['dried_wf_g']) print("\n--- Empirical Water Solubility (Absolute) ---") print(f" Dry WF Input: {dry_wf_in:.4f} g") print(f" S_water at T_low: {s_low*100:.1f}% by mass ({water_extracted_g:.4f} g extracted)") print(f" S_water at T_high: {s_high*100:.1f}% by mass ({max(0.0, water_trapped_g):.4f} g trapped)") # ========================================== # 2. VOLUMETRIC SALT PARTITIONING (mg/L) # ========================================== k_low, k_high = 0.0, 0.0 true_reject_mass_g = inputs['feed_water_g'] - water_extracted_g if outputs['reject_tds_mg_kg'] is not None and outputs['yield_tds_mg_kg'] is not None: # Calculate absolute salt (mg) feed_salt_mg = inputs['feed_water_g'] * (inputs['feed_tds_mg_kg'] / 1000.0) reject_salt_mg = true_reject_mass_g * (outputs['reject_tds_mg_kg'] / 1000.0) yield_salt_mg = outputs['yield_water_g'] * (outputs['yield_tds_mg_kg'] / 1000.0) # --- Density & Volume Conversions --- wf_density_kg_L = inputs.get('wf_density_kg_L', 0.80) # Aqueous Volumes (using empirical density approximation) reject_density = 1.0 + (outputs['reject_tds_mg_kg'] / 1_000_000.0) v_aq_reject_L = (true_reject_mass_g / 1000.0) / reject_density yield_density = 1.0 + (outputs['yield_tds_mg_kg'] / 1_000_000.0) v_aq_yield_L = (outputs['yield_water_g'] / 1000.0) / yield_density # Organic Volumes (Additive assumption: V_total = V_dry_wf + V_water) v_dry_wf_L = (dry_wf_in / 1000.0) / wf_density_kg_L v_org_hydrated_L = v_dry_wf_L + (water_extracted_g / 1000.0) / 1.0 # water density ~1.0 kg/L v_org_dried_L = v_dry_wf_L + (water_trapped_g / 1000.0) / 1.0 # --- K_low (Volumetric) --- salt_transferred_mg = feed_salt_mg - reject_salt_mg c_org_hydrated_vol = salt_transferred_mg / v_org_hydrated_L if v_org_hydrated_L > 0 else 0 c_aq_reject_vol = reject_salt_mg / v_aq_reject_L if v_aq_reject_L > 0 else 0 k_low = c_org_hydrated_vol / c_aq_reject_vol if c_aq_reject_vol > 0 else 0 # --- K_high (Volumetric) --- salt_trapped_mg = salt_transferred_mg - yield_salt_mg c_org_dried_vol = salt_trapped_mg / v_org_dried_L if v_org_dried_L > 0 else 0 c_aq_yield_vol = yield_salt_mg / v_aq_yield_L if v_aq_yield_L > 0 else 0 k_high = c_org_dried_vol / c_aq_yield_vol if c_aq_yield_vol > 0 else 0 print("\n--- Empirical Salt Partitioning (Volumetric mg/L) ---") print(f" Calculated K_salt at T_low: {k_low:.5f}") print(f" Calculated K_salt at T_high: {k_high:.5f}") # ========================================== # 3. MASS CLOSURE & DATA INTEGRITY REPORT # ========================================== print("\n--- Mass Closure & Data Integrity Report ---") stage2_input = outputs['hydrated_wf_g'] stage2_output = outputs['yield_water_g'] + outputs['dried_wf_g'] stage2_diff = stage2_input - stage2_output if abs(stage2_diff) < 0.001: print(" ✅ Stage 2 Mass Balance: Perfect closure.") else: print(f" ❌ Stage 2 Mass Balance Error: {stage2_diff:+.4f} g missing during separation.") if water_trapped_g < -0.01: organic_loss = abs(water_trapped_g) loss_pct = (organic_loss / dry_wf_in) * 100 print(f" ⚠️ Organic Phase Loss: {organic_loss:.4f} g ({loss_pct:.1f}%) of the working fluid was lost.") else: print(" ✅ Organic Phase Mass: No obvious solvent loss detected.") reported_reject = outputs.get('reject_water_g') if reported_reject is not None: table_artifact = inputs['feed_water_g'] - outputs['yield_water_g'] if abs(reported_reject - table_artifact) < 0.001: print(f" ⚠️ Reject Mass Artifact: The reported Reject water mass appears to be a spreadsheet calculation.") print(f" -> Script automatically corrected Reject to: {true_reject_mass_g:.4f} g") # ========================================== # 4. EXPORT TO SPECIFIC YAML # ========================================== acid_name = run_data.get('acid_used', 'Unknown_Acid') export_data = { "solvent_properties": { "name": f"WF2 + {acid_name}", "water_solubility": { "t_low_mass_fraction": float(round(s_low, 4)), "t_high_mass_fraction": float(round(s_high, 4)) }, "ion_partition_coefficients": { "t_low": {"default": float(round(k_low, 6))}, "t_high": {"default": float(round(k_high, 6))} } } } output_filename = f"empirical_wf_{experiment_name}.yaml" with open(output_filename, 'w') as file: yaml.safe_dump(export_data, file, sort_keys=False) def stage_1_residuals(x, species_list, feed_masses_g, solvent_mass_g, params): """ The physics engine core. Evaluates the mass balances and equilibrium equations. x = array of guesses for the mass of each species in the organic phase. """ k_dict = params['solvent_properties']['ion_partition_coefficients']['t_low'] default_k = k_dict['default'] water_solubility = params['solvent_properties']['water_solubility']['t_low_mass_fraction'] # 1. Distribute the mass based on Python's current guesses (x) org_masses = {species: x[i] for i, species in enumerate(species_list)} aq_masses = {species: feed_masses_g[species] - x[i] for i, species in enumerate(species_list)} # 2. Calculate dynamic totals and volumes total_org_mass = solvent_mass_g + sum(org_masses.values()) # Convert density from kg/L to g/L for the math solvent_density_g_L = params['solvent_properties'].get('density_kg_L', 0.8) * 1000.0 org_vol_L = total_org_mass / solvent_density_g_L total_aq_mass = sum(aq_masses.values()) aq_ion_mass = total_aq_mass - aq_masses['H2O'] aq_density = calculate_brine_density(aq_ion_mass, total_aq_mass) aq_vol_L = total_aq_mass / (aq_density * 1000.0) # 3. Calculate the residuals (how far off the physics are from perfect equilibrium) residuals = [] for species in species_list: m_org = org_masses[species] m_aq = aq_masses[species] if species == 'H2O': # Is the solvent saturated with water? actual_fraction = m_org / total_org_mass residuals.append(water_solubility - actual_fraction) else: # Did the ions partition correctly? c_org = m_org / org_vol_L c_aq = m_aq / aq_vol_L if aq_vol_L > 0 else 0 k_val = k_dict.get(species, default_k) residuals.append((k_val * c_aq) - c_org) return residuals def stage_2_residuals(x, species_list, loaded_org_masses_g, solvent_mass_g, params): """ Evaluates the recovery stage at T_high. x = array of guesses for the mass of each species STAYING in the organic phase. """ k_dict = params['solvent_properties']['ion_partition_coefficients']['t_high'] default_k = k_dict['default'] water_solubility = params['solvent_properties']['water_solubility']['t_high_mass_fraction'] # 1. Distribute the mass based on guesses # x is what stays in the organic phase. The rest is expelled as our aqueous yield. org_masses = {species: x[i] for i, species in enumerate(species_list)} aq_masses = {species: loaded_org_masses_g[species] - x[i] for i, species in enumerate(species_list)} # 2. Calculate dynamic totals and volumes total_org_mass = solvent_mass_g + sum(org_masses.values()) # Convert density from kg/L to g/L for the math solvent_density_g_L = params['solvent_properties'].get('density_kg_L', 0.8) * 1000.0 org_vol_L = total_org_mass / solvent_density_g_L total_aq_mass = sum(aq_masses.values()) # Prevent division by zero if the solver guesses 0 aqueous mass if total_aq_mass > 0: aq_ion_mass = total_aq_mass - aq_masses['H2O'] aq_density = calculate_brine_density(aq_ion_mass, total_aq_mass) aq_vol_L = total_aq_mass / (aq_density * 1000.0) else: aq_vol_L = 0.0001 # tiny buffer # 3. Calculate the residuals residuals = [] for species in species_list: m_org = org_masses[species] m_aq = aq_masses[species] if species == 'H2O': actual_fraction = m_org / total_org_mass if total_org_mass > 0 else 0 residuals.append(water_solubility - actual_fraction) else: c_org = m_org / org_vol_L if org_vol_L > 0 else 0 c_aq = m_aq / aq_vol_L if aq_vol_L > 0 else 0 k_val = k_dict.get(species, default_k) residuals.append((k_val * c_aq) - c_org) return residuals def convert_mass_to_concentration(masses_g): """ Converts an array of absolute masses (grams) back into Total Volume (L) and Concentrations (mg/L). """ total_mass_g = sum(masses_g.values()) # Safety check: if the phase is virtually empty, return zeros if total_mass_g < 0.001: return 0.0, {} ion_mass_g = total_mass_g - masses_g.get('H2O', 0.0) # Calculate Density and Volume density_kg_L = calculate_brine_density(ion_mass_g, total_mass_g) volume_L = (total_mass_g / 1000.0) / density_kg_L # Calculate mg/L for each ion concentrations_mg_L = {} for species, mass in masses_g.items(): if species != 'H2O' and mass > 0.0001: # (grams * 1000) / Liters = mg/L concentrations_mg_L[species] = (mass * 1000.0) / volume_L return volume_L, concentrations_mg_L def process_all_empirical_runs(experimental_yaml_path="experimental_runs.yaml"): """ Loops through all experiments in the YAML file and processes them. Structured to easily drop in a hashing/caching mechanism later. """ print(f"\n📂 Loading batch experimental data from {experimental_yaml_path}...") data = load_yaml(experimental_yaml_path) experiments = data.get('experiments', {}) # ========================================== # FUTURE CACHE SETUP # ========================================== # hash_file = ".empirical_cache.yaml" # previous_hashes = load_yaml(hash_file) if os.path.exists(hash_file) else {} # current_hashes = {} processed_count = 0 for experiment_name, run_data in experiments.items(): # ========================================== # FUTURE CACHE CHECK # ========================================== # run_hash = hashlib.sha256(str(run_data).encode('utf-8')).hexdigest() # current_hashes[experiment_name] = run_hash # if previous_hashes.get(experiment_name) == run_hash: # print(f"⏩ Skipping {experiment_name} (No changes detected).") # continue # Pass the dictionary directly to the math engine run_empirical_mode(experiment_name, run_data) processed_count += 1 # ========================================== # FUTURE CACHE SAVE # ========================================== # with open(hash_file, 'w') as f: # yaml.safe_dump(current_hashes, f) print(f"\n🎉 Finished processing {processed_count} experimental runs!") def main(): print("🌊 Starting TS-LLE Solver...") # 1. Load Files water_file = "water_sources/permian_brine.yaml" params_file = "tsse_parameters.yaml" feed_water = load_yaml(water_file) params = load_yaml(params_file) # ========================================== # NEW: EMPIRICAL OVERRIDE LOGIC # ========================================== empirical_file = params['process_settings'].get('empirical_solvent_file') if empirical_file and os.path.exists(empirical_file): print(f"🔄 Overwriting default solvent thermodynamics with: {empirical_file}") emp_data = load_yaml(empirical_file) emp_props = emp_data.get('solvent_properties', {}) # Safely overwrite only the keys generated by the empirical parser if 'name' in emp_props: params['solvent_properties']['name'] = emp_props['name'] if 'water_solubility' in emp_props: params['solvent_properties']['water_solubility'] = emp_props['water_solubility'] if 'ion_partition_coefficients' in emp_props: params['solvent_properties']['ion_partition_coefficients'] = emp_props['ion_partition_coefficients'] elif empirical_file: print(f"⚠️ WARNING: Specified empirical file '{empirical_file}' not found. Using defaults.") # 2. Extract Base Variables mode = params['process_settings'].get('mode', 'forward') feed_volume = params['process_settings']['feed_volume_L'] solvent_volume_L = params['process_settings']['solvent_volume_L'] ions_dict = feed_water.get('ions', {}) # Dynamically calculate solvent mass based on YAML density (which is safely preserved!) density_kg_L = params['solvent_properties'].get('density_kg_L', 0.8) solvent_mass_g = solvent_volume_L * (density_kg_L * 1000.0) # 3. Route to the appropriate engine if mode == "reverse": target_reject_L = params['process_settings'].get('target_reject_volume_L', feed_volume / 2.0) yield_volume_L = feed_volume - target_reject_L target_yield_g = yield_volume_L * 1000.0 s_low_max = params['solvent_properties'].get('max_water_uptake_frac', 0.833) run_reverse_mode(target_yield_g, solvent_mass_g, s_low_max) elif mode == "forward": run_forward_mode(params, ions_dict, feed_volume, solvent_mass_g) elif mode == "empirical": process_all_empirical_runs("experimental_runs.yaml") else: print(f"❌ Unknown mode selected in config: {mode}") if __name__ == "__main__": main()