import warnings
import time
from typing import List, Union, Dict, Any, Optional, Tuple
from AFL.automation.mixing.MassBalance import MassBalance
from AFL.automation.mixing.MassBalanceDriver import MassBalanceDriver
from AFL.automation.prepare.OT2HTTPDriver import OT2HTTPDriver
from AFL.automation.shared.utilities import listify
from AFL.automation.shared.units import to_quantity, is_volume
from AFL.automation.mixing.Solution import Solution
from AFL.automation.mixing.PipetteAction import PipetteAction
[docs]
class OT2Prepare(OT2HTTPDriver, MassBalanceDriver):
defaults = {
'prep_targets': [],
'prepare_volume': '900 ul',
'catch_volume': '900 ul',
'deck': {},
'stocks': [],
'stock_mix_order': [],
'fixed_compositions': {},
'stock_locations': {}, # Maps stock names to deck positions: {'stockH2O': '3A2'}
'stock_transfer_params': {}, # Per-stock transfer parameters: {'stockH2O': {'mix_after': True}}
'catch_protocol': {}, # PipetteAction-formatted dict for catch transfer parameters
}
[docs]
def __init__(self, overrides=None):
# Initialize both parent classes
OT2HTTPDriver.__init__(self, overrides=overrides)
MassBalanceDriver.__init__(self, overrides=overrides)
# Override the name set by both parents
self.name = 'OT2Prepare'
# Update filepath to match the new name
self.filepath = self.path / (self.name + '.config.json')
# Replace config with optimized settings for OT2Prepare
# Inherits optimized settings from MassBalanceDriver, but we can further customize
# Note: PersistentConfig will automatically load existing values from disk
from AFL.automation.shared.PersistentConfig import PersistentConfig
self.config = PersistentConfig(
path=self.filepath,
defaults=self.gather_defaults(),
overrides=overrides,
max_history=100, # Reduced from default 10000
max_history_size_mb=50, # Limit file size to 50MB
write_debounce_seconds=0.5, # Batch rapid operations
compact_json=True, # Use compact JSON for large files
)
# Initialize additional attributes
self.stocks = []
self.targets = []
self.last_target_location = None
self.useful_links['View Deck'] = '/visualize_deck'
self.process_stocks()
[docs]
def status(self):
"""
Get the status of the OT2Prepare driver.
"""
# Get status from OT2HTTPDriver
ot2_status = OT2HTTPDriver.status(self)
# Add our own status information
status = []
status.append(f'Stocks: {len(self.stocks)} configured')
status.append(f'Stock locations: {self.config["stock_locations"]}')
status.append(f'{len(self.config["prep_targets"])} preparation targets available')
# Combine status information
return status + ot2_status
[docs]
def is_feasible(self, targets: dict | list[dict]) -> list[dict | None]:
"""
Check if the target composition(s) is/are feasible for preparation using mass balance.
If feasible, returns the balanced target solution dictionary. Otherwise, returns None.
This implementation creates a local MassBalance instance for each feasibility check
to avoid modifying the driver's state.
Parameters
----------
targets : Union[dict, List[dict]]
Either a single target dictionary or a list of target dictionaries.
Returns
-------
List[Union[dict, None]]
A list containing the balanced target dictionary for each feasible target,
or None for infeasible targets.
"""
targets_to_check = listify(targets)
# Process stocks from the driver
self.process_stocks()
# Get the minimum volume configuration
minimum_volume = self.config.get('minimum_volume', '100 ul')
results = []
for target in targets_to_check:
try:
# Create a local MassBalance instance
mb = MassBalance(minimum_volume=minimum_volume)
# Configure the same stocks as in the driver
for stock in self.stocks:
mb.stocks.append(stock)
# Apply any fixed compositions from config
target_with_fixed = self.apply_fixed_comps(target.copy())
# Create a Solution from the target and add it to the MassBalance instance
from AFL.automation.mixing.Solution import Solution
target_solution = Solution(**target_with_fixed)
mb.targets.append(target_solution)
# Calculate mass balance
mb.balance(tol=self.config.get('tol', 1e-3))
# Check if balance was successful for this target
if (mb.balanced and
len(mb.balanced) > 0 and
mb.balanced[0].get('balanced_target') is not None):
results.append(mb.balanced[0]['balanced_target'].to_dict())
else:
results.append(None)
except Exception as e:
# If an exception occurs, indicate failure
warnings.warn(f"Exception during feasibility check for target {target.get('name', 'Unnamed')}: {str(e)}", stacklevel=2)
results.append(None)
return results
[docs]
def apply_fixed_comps(self, target: dict) -> dict:
"""
Apply fixed compositions to a target dictionary without overwriting existing values.
Parameters
----------
target : dict
The target solution dictionary
Returns
-------
dict
A new target dictionary with fixed compositions applied
"""
# Create a copy to avoid modifying the original
result = target.copy()
# Get fixed compositions from config
fixed_comps = self.config.get('fixed_compositions', {})
if not fixed_comps:
return result
# For each component property type that might exist in the target
for prop_type in ['masses', 'volumes', 'concentrations', 'mass_fractions']:
# Initialize property dictionaries if they don't exist
if prop_type not in result:
result[prop_type] = {}
# If this property exists in fixed compositions
if prop_type in fixed_comps:
# Add each component from fixed compositions that doesn't already exist
for comp_name, comp_value in fixed_comps[prop_type].items():
if comp_name not in result[prop_type]:
result[prop_type][comp_name] = comp_value
# Handle simpler properties that might not be dictionaries
for prop in ['total_mass', 'total_volume', 'name', 'location']:
if prop in fixed_comps and prop not in result:
result[prop] = fixed_comps[prop]
# Handle solutes list
if 'solutes' in fixed_comps:
if 'solutes' not in result:
result['solutes'] = fixed_comps['solutes'].copy()
else:
# Add any solutes that aren't already in the list
for solute in fixed_comps['solutes']:
if solute not in result['solutes']:
result['solutes'].append(solute)
return result
[docs]
def prepare(self, target: dict, dest: str | None = None) -> tuple[dict, str] | tuple[None, None]:
"""Prepare the target solution. The dest argument is currently not used by this implementation."""
# Apply fixed compositions without overwriting existing values
target = self.apply_fixed_comps(target)
# Check if the target is feasible before attempting preparation
feasibility_results = self.is_feasible(target)
if not feasibility_results or feasibility_results[0] is None:
warnings.warn(f'Target composition {target.get("name", "Unnamed target")} is not feasible based on mass balance calculations', stacklevel=2)
return None, None
balanced_target_dict_from_feasible = feasibility_results[0]
self.reset_targets()
# We need to re-add the original target, not the dict from is_feasible
self.add_target(target)
self.balance()
if not self.balanced or not self.balanced[0].get('balanced_target'):
warnings.warn(f'No suitable mass balance found for target: {target.get("name", "Unnamed target")}',stacklevel=2)
return None, None
# This is the Solution object containing the protocol
balanced_target_solution_object = self.balanced[0]['balanced_target']
# Configure the destination for the preparation
if not self.config.get('prep_targets'):
raise ValueError("No preparation targets configured. Cannot select a destination target.")
# Pop a location from the preparation targets list
if dest is None:
# need to pop and then resend the locations list so that the persistant config triggers a write
prep_targets = self.config['prep_targets']
destination = prep_targets.pop(0)
self.config['prep_targets'] = prep_targets
else:
destination = dest
# Execute the protocol using OT2HTTPDriver
if not hasattr(balanced_target_solution_object, 'protocol') or not balanced_target_solution_object.protocol:
raise ValueError("No protocol generated for the target solution")
# Reorder protocol based on stock_mix_order if specified
protocol = self.reorder_protocol(balanced_target_solution_object.protocol)
# Execute each step in the protocol
for step in protocol:
# Get source and destination
# Note: source is already a deck location (e.g., '3A2')
source = step.source
dest = destination
volume_ul = step.volume # Volume is already in μL
# Get stock name from deck location for transfer parameters
stock_name = self.config.get('deck', {}).get(source)
if stock_name is None:
raise ValueError(f"No stock name found for deck location: {source}")
# Get stock-specific transfer parameters
transfer_params = self.get_transfer_params(stock_name)
# Execute the transfer
try:
self.transfer(
source=source,
dest=dest,
volume=volume_ul,
**transfer_params
)
except Exception as e:
warnings.warn(f"Transfer failed from {source} to {dest}: {str(e)}", stacklevel=2)
return None, None
# Store the last target location for catch transfer
self.last_target_location = destination
# Return the balanced target and destination
result_dict = balanced_target_solution_object.to_dict()
# Add total_volume to the result dictionary
if hasattr(balanced_target_solution_object, 'volume') and balanced_target_solution_object.volume is not None:
result_dict['total_volume'] = f"{balanced_target_solution_object.volume.to('ul').magnitude} ul"
return result_dict, destination
[docs]
def process_stocks(self):
"""
Process stocks and update deck config with inverse of stock_locations.
"""
# Call parent method to process stocks and update stock_locations
MassBalanceDriver.process_stocks(self)
# Populate deck config with inverse of stock_locations
self._update_deck_config()
def _update_deck_config(self):
"""
Update the deck config with the inverse of stock_locations.
This creates a mapping from deck locations to stock names.
"""
deck_config = {}
stock_locations = self.config.get('stock_locations', {})
for stock_name, deck_location in stock_locations.items():
deck_config[deck_location] = stock_name
self.config['deck'] = deck_config
[docs]
def get_transfer_params(self, stock_name):
"""
Get the transfer parameters for a specific stock solution.
Parameters
----------
stock_name : str
Name of the stock solution
Returns
-------
dict
Dictionary of transfer parameters to pass to transfer()
"""
# Get stock-specific parameters if available
stock_params = self.config.get('stock_transfer_params', {}).get(stock_name, {})
# Get default parameters
default_params = self.config.get('stock_transfer_params', {}).get('default', {})
# Combine default and stock-specific parameters, with stock-specific taking precedence
params = default_params.copy()
params.update(stock_params)
return params
[docs]
def reorder_protocol(self, protocol):
"""
Reorder the protocol based on stock_mix_order if specified
Parameters
----------
protocol : list
List of PipetteAction objects
Returns
-------
list
Reordered list of PipetteAction objects
"""
# If no stock_mix_order is specified, return original protocol
stock_mix_order = self.config.get('stock_mix_order', [])
if not stock_mix_order:
return protocol
# Group protocol steps by source
steps_by_source = {}
for step in protocol:
if step.source not in steps_by_source:
steps_by_source[step.source] = []
steps_by_source[step.source].append(step)
# Build reordered protocol based on stock_mix_order
reordered = []
for stock_name in stock_mix_order:
if stock_name in steps_by_source:
reordered.extend(steps_by_source[stock_name])
del steps_by_source[stock_name]
# Add any remaining steps that weren't in stock_mix_order
for steps in steps_by_source.values():
reordered.extend(steps)
return reordered
[docs]
def transfer_to_catch(self, source=None, dest=None, **kwargs):
"""
Transfer a prepared sample to the catch/loader location using catch protocol settings.
Parameters
----------
source : str, optional
Source location (well) of the prepared sample. If None, uses the last prepared target location.
dest : str, optional
Destination location (well). If None, must be specified in catch_protocol config.
**kwargs
Additional transfer parameters that override catch_protocol settings.
Returns
-------
str or None
UUID of the transfer task if successful, None otherwise
"""
# Get catch protocol parameters from config
catch_params = self.config.get('catch_protocol', {}).copy()
# Determine source location
if source is None:
if self.last_target_location is None:
raise ValueError("No source specified and no last target location available. Call prepare() first or specify source.")
source = self.last_target_location
kwargs['source'] = source
# Handle destination
if dest is not None:
kwargs['dest'] = dest
# Merge with kwargs (kwargs take precedence)
catch_params.update(kwargs)
if 'dest' not in catch_params:
raise ValueError("Destination 'dest' must be specified in catch_protocol config or as an argument.")
# Execute the transfer
try:
self.transfer(**catch_params)
except Exception as e:
dest_val = catch_params.get('dest', 'unknown')
warnings.warn(f"Transfer to catch failed from {source} to {dest_val} using {catch_params}: {str(e)}", stacklevel=2)
raise
[docs]
def reset(self):
"""Reset the driver state/configuration."""
# Placeholder: implement reset logic
self.reset_targets()
self.reset_stocks()
_DEFAULT_PORT=5002
if __name__ == '__main__':
from AFL.automation.shared.launcher import *
