Coverage for backend/ahuora-builder/src/ahuora_builder/custom/watertap/reverse_osmosis_0d.py: 86%
166 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-07-22 05:22 +0000
« prev ^ index » next coverage.py v7.10.7, created at 2026-07-22 05:22 +0000
1from idaes.core import declare_process_block_class
2from watertap.unit_models.reverse_osmosis_0D import ReverseOsmosisData
3from pyomo.common.collections import ComponentSet
4from pyomo.core.base.component import Component
5from pyomo.core.expr.visitor import identify_variables
6from pyomo.environ import Constraint, Expression, value
7from pyomo.repn import generate_standard_repn
8from idaes.core.util.model_statistics import degrees_of_freedom
9from idaes.core.util import scaling as iscale
10from watertap.core.membrane_channel_base import TransportModel
11from ahuora_builder.methods.scaling_suffix import sanitize_scaling_suffix
15@declare_process_block_class("ReverseOsmosis0D")
16class ReverseOsmosis0DData(ReverseOsmosisData):
17 """
18 Ahuora wrapper around WaterTAP's ReverseOsmosis0D unit model.
19 """
21 def build(self):
22 super().build()
23 time = self.flowsheet().time
25 self.inlet_reynolds_number = Expression(
26 time,
27 rule=lambda b, t: b.feed_side.N_Re[
28 t, b.feed_side.length_domain.first()
29 ],
30 doc="Feed-side Reynolds number at the channel inlet",
31 )
32 self.volumetric_recovery = Expression(
33 time,
34 rule=lambda b, t: b.recovery_vol_phase[t, "Liq"],
35 doc="Liquid-phase volumetric recovery",
36 )
37 self.tds_mass_recovery = Expression(
38 time,
39 rule=lambda b, t: b.recovery_mass_phase_comp[t, "Liq", "TDS"],
40 doc="TDS mass recovery",
41 )
42 self.tds_observed_rejection = Expression(
43 time,
44 rule=lambda b, t: b.rejection_phase_comp[t, "Liq", "TDS"],
45 doc="Observed TDS rejection",
46 )
47 self.average_water_mass_flux = Expression(
48 time,
49 rule=lambda b, t: b.flux_mass_phase_comp_avg[t, "Liq", "H2O"],
50 doc="Average water mass flux through the membrane",
51 )
53 def initialize_build(self, *args, **kwargs):
54 iscale.calculate_scaling_factors(self)
55 sanitize_scaling_suffix(self)
56 return super().initialize_build(*args, **kwargs)
58 def diagnose(self) -> list[tuple[Component, str]]:
59 """
60 Return user-facing hints for common reverse osmosis specification issues.
62 RO needs a pressure driving force, valid membrane transport parameters,
63 and exactly enough geometry specifications to satisfy ``area = length *
64 width``. Recovery, rejection, flux, and over-pressure ratio are calculated
65 outputs in this configuration, so fixing them can over-specify the unit.
66 """
67 problems: list[tuple[Component, str]] = []
69 feed_in = self.feed_side.properties_in[0]
70 permeate = self.mixed_permeate[0]
71 feed_pressure = _safe_value(feed_in.pressure)
72 permeate_pressure = _safe_value(permeate.pressure)
73 if feed_pressure is not None and permeate_pressure is not None: 73 ↛ 92line 73 didn't jump to line 92 because the condition on line 73 was always true
74 if permeate_pressure >= feed_pressure:
75 problems.append(
76 (
77 permeate.pressure,
78 "The reverse osmosis permeate pressure is not lower than "
79 "feed pressure. RO needs feed pressure above permeate "
80 "pressure to drive water through the membrane.",
81 )
82 )
83 elif feed_pressure - permeate_pressure < 1e5: 83 ↛ 84line 83 didn't jump to line 84 because the condition on line 83 was never true
84 problems.append(
85 (
86 feed_in.pressure,
87 "The reverse osmosis pressure difference is very small. "
88 "This may lead to little or no permeate production.",
89 )
90 )
92 for component in self.A_comp.values():
93 water_permeability = _safe_value(component)
94 if water_permeability is not None and water_permeability <= 0:
95 problems.append(
96 (
97 component,
98 "Reverse osmosis water permeability must be greater than 0.",
99 )
100 )
102 for component in self.B_comp.values():
103 salt_permeability = _safe_value(component)
104 if salt_permeability is not None and salt_permeability < 0:
105 problems.append(
106 (
107 component,
108 "Reverse osmosis salt permeability must be greater than or "
109 "equal to 0.",
110 )
111 )
113 geometry_specs = [
114 _is_specified(self, self.area),
115 _is_specified(self, self.length),
116 _is_specified(self, self.width),
117 ]
118 if sum(geometry_specs) < 2:
119 problems.append(
120 (
121 self.area,
122 "Reverse osmosis membrane geometry is under-specified. Specify "
123 "two of membrane area, length, and width.",
124 )
125 )
126 elif all(geometry_specs):
127 problems.append(
128 (
129 self.width,
130 "Reverse osmosis membrane area, length, and width are all fixed. "
131 "Because area = length * width, fixing all three can "
132 "over-specify the unit.",
133 )
134 )
136 for component, label in (
137 (self.area, "membrane area"),
138 (self.length, "membrane length"),
139 (self.width, "membrane width"),
140 ):
141 component_value = _safe_value(component)
142 if component_value is not None and component_value <= 0:
143 problems.append(
144 (
145 component,
146 f"Reverse osmosis {label} must be greater than 0.",
147 )
148 )
150 spacer_porosity = _safe_value(self.feed_side.spacer_porosity)
151 if spacer_porosity is not None and not 0 < spacer_porosity < 1:
152 problems.append(
153 (
154 self.feed_side.spacer_porosity,
155 "Reverse osmosis feed spacer porosity must be between 0 and 1.",
156 )
157 )
159 channel_height = _safe_value(self.feed_side.channel_height)
160 if channel_height is not None and channel_height <= 0:
161 problems.append(
162 (
163 self.feed_side.channel_height,
164 "Reverse osmosis feed channel height must be greater than 0.",
165 )
166 )
168 for component, label in (
169 (self.recovery_vol_phase[0, "Liq"], "volumetric recovery"),
170 (
171 self.recovery_mass_phase_comp[0, "Liq", "TDS"],
172 "TDS mass recovery",
173 ),
174 (
175 self.rejection_phase_comp[0, "Liq", "TDS"],
176 "observed solute rejection",
177 ),
178 ):
179 component_value = _safe_value(component)
180 if component_value is not None and not 0 <= component_value <= 1:
181 problems.append(
182 (
183 component,
184 f"Reverse osmosis {label} should be between 0 and 1.",
185 )
186 )
188 calculated_outputs = (
189 (
190 self.recovery_vol_phase[0, "Liq"],
191 "Volumetric recovery is calculated from the RO mass balances.",
192 ),
193 (
194 self.recovery_mass_phase_comp[0, "Liq", "TDS"],
195 "Mass recovery is calculated from the RO component balances.",
196 ),
197 (
198 self.rejection_phase_comp[0, "Liq", "TDS"],
199 "Observed solute rejection is calculated from feed and permeate "
200 "composition.",
201 ),
202 (
203 self.flux_mass_phase_comp_avg[0, "Liq", "H2O"],
204 "Average mass flux is calculated from membrane transport equations.",
205 ),
206 (
207 self.over_pressure_ratio[0],
208 "Over-pressure ratio is calculated from feed pressure and osmotic "
209 "pressure.",
210 ),
211 )
212 model_dof = _safe_degrees_of_freedom(self)
213 for component, message in calculated_outputs:
214 if model_dof is not None and model_dof < 0 and _is_specified(self, component):
215 problems.append(
216 (
217 component,
218 f"{message} Fixing it can over-specify the reverse osmosis unit.",
219 )
220 )
222 return problems
224 def calculate_scaling_factors(self):
225 # these variables should have user input, if not there will be a warning
226 if iscale.get_scaling_factor(self.area) is None:
227 sf = iscale.get_scaling_factor(self.area, default=10, warning=True)
228 iscale.set_scaling_factor(self.area, sf)
230 if iscale.get_scaling_factor(self.A_comp) is None:
231 iscale.set_scaling_factor(self.A_comp, 1e12)
233 if iscale.get_scaling_factor(self.B_comp) is None:
234 iscale.set_scaling_factor(self.B_comp, 1e8)
236 if iscale.get_scaling_factor(self.recovery_vol_phase) is None:
237 iscale.set_scaling_factor(self.recovery_vol_phase, 1)
239 if self.config.transport_model == TransportModel.SKK: 239 ↛ 240line 239 didn't jump to line 240 because the condition on line 239 was never true
240 if iscale.get_scaling_factor(self.alpha) is None:
241 iscale.set_scaling_factor(self.alpha, 1e-8)
243 if iscale.get_scaling_factor(self.reflect_coeff) is None:
244 iscale.set_scaling_factor(self.reflect_coeff, 1)
246 for (t, p, j), v in self.recovery_mass_phase_comp.items():
247 if j in self.config.property_package.solvent_set:
248 sf = 1
249 elif j in self.config.property_package.solute_set: 249 ↛ 251line 249 didn't jump to line 251 because the condition on line 249 was always true
250 sf = 100
251 if iscale.get_scaling_factor(v) is None:
252 iscale.set_scaling_factor(v, sf)
254 for v in self.rejection_phase_comp.values():
255 if iscale.get_scaling_factor(v) is None:
256 iscale.set_scaling_factor(v, 1)
258 if not hasattr(self, "_permeate_scaled_properties"):
259 self._permeate_scaled_properties = ComponentSet()
261 for sb in (self.permeate_side, self.mixed_permeate):
262 for blk in sb.values():
263 for j in self.config.property_package.solute_set:
264 self._rescale_permeate_variable(blk.flow_mass_phase_comp["Liq", j])
265 if blk.is_property_constructed("mass_frac_phase_comp"): 265 ↛ 269line 265 didn't jump to line 269 because the condition on line 265 was always true
266 self._rescale_permeate_variable(
267 blk.mass_frac_phase_comp["Liq", j]
268 )
269 if blk.is_property_constructed("conc_mass_phase_comp"): 269 ↛ 273line 269 didn't jump to line 273 because the condition on line 269 was always true
270 self._rescale_permeate_variable(
271 blk.conc_mass_phase_comp["Liq", j]
272 )
273 if blk.is_property_constructed("mole_frac_phase_comp"):
274 self._rescale_permeate_variable(blk.mole_frac_phase_comp["Liq",j])
275 if blk.is_property_constructed("molality_phase_comp"):
276 self._rescale_permeate_variable(
277 blk.molality_phase_comp["Liq", j]
278 )
279 if blk.is_property_constructed("pressure_osm_phase"):
280 self._rescale_permeate_variable(blk.pressure_osm_phase["Liq"])
282 for (t, x, p, j), v in self.flux_mass_phase_comp.items():
283 if iscale.get_scaling_factor(v) is None:
284 comp = self.config.property_package.get_component(j)
285 if comp.is_solvent(): # scaling based on solvent flux equation
286 sf = (
287 iscale.get_scaling_factor(self.A_comp[t, j])
288 * iscale.get_scaling_factor(self.dens_solvent)
289 * iscale.get_scaling_factor(
290 self.feed_side.properties[t, x].pressure
291 )
292 )
293 iscale.set_scaling_factor(v, sf)
294 elif comp.is_solute(): # scaling based on solute flux equation 294 ↛ 282line 294 didn't jump to line 282 because the condition on line 294 was always true
295 sf = iscale.get_scaling_factor(
296 self.B_comp[t, j]
297 ) * iscale.get_scaling_factor(
298 self.feed_side.properties[t, x].conc_mass_phase_comp[p, j]
299 )
300 iscale.set_scaling_factor(v, sf)
302 @staticmethod
303 def ahuora_metadata():
304 return _reverse_osmosis_0d_metadata()
307def _safe_value(component) -> float | None:
308 try:
309 return value(component, exception=False)
310 except Exception:
311 return None
314def _safe_degrees_of_freedom(block) -> int | None:
315 try:
316 return degrees_of_freedom(block.parent_block())
317 except Exception:
318 try:
319 return degrees_of_freedom(block)
320 except Exception:
321 return None
324def _is_specified(block, component) -> bool:
325 return getattr(component, "fixed", False) or _has_external_equality_spec(
326 block, component
327 )
330def _has_external_equality_spec(block, component) -> bool:
331 target_vars = ComponentSet(identify_variables(component, include_fixed=False))
332 if not target_vars and hasattr(component, "is_variable_type"): 332 ↛ 333line 332 didn't jump to line 333 because the condition on line 332 was never true
333 target_vars = ComponentSet([component])
334 if not target_vars: 334 ↛ 335line 334 didn't jump to line 335 because the condition on line 334 was never true
335 return False
337 for constraint in block.parent_block().component_data_objects(
338 Constraint, active=True, descend_into=True
339 ):
340 if _is_descendant_of(constraint.parent_block(), block):
341 continue
342 if not constraint.equality: 342 ↛ 343line 342 didn't jump to line 343 because the condition on line 342 was never true
343 continue
345 repn = generate_standard_repn(constraint.body)
346 if repn.is_linear() and len(repn.linear_vars) == 1: 346 ↛ 347line 346 didn't jump to line 347 because the condition on line 346 was never true
347 if repn.linear_vars[0] in target_vars:
348 return True
350 constraint_vars = ComponentSet(
351 identify_variables(constraint.body, include_fixed=False)
352 )
353 if constraint_vars and constraint_vars == target_vars:
354 return True
356 return False
359def _is_descendant_of(child_block, parent_block) -> bool:
360 block = child_block
361 while block is not None:
362 if block is parent_block:
363 return True
364 block = block.parent_block()
365 return False
368def _reverse_osmosis_0d_metadata():
369 from ahuora_unit_ops.json_config import (
370 JsonAdapterArgConfig,
371 JsonFrontendConfig,
372 JsonGraphicObjectConfig,
373 JsonIdaesAdapterConfig,
374 JsonPortConfig,
375 JsonPropertyConfig,
376 JsonPropertySetGroupConfig,
377 JsonUnitOpConfig,
378 )
380 return JsonUnitOpConfig(
381 key='reverse_osmosis_0d',
382 objectType='reverse_osmosis_0d',
383 enumMember='ReverseOsmosis0D',
384 displayType='Reverse Osmosis (0D)',
385 displayName='Reverse Osmosis (0D)',
386 categoryPath=['chemical', 'separation'],
387 ports={
388 'inlet': JsonPortConfig(
389 displayName='Inlet',
390 type='inlet',
391 streamType='stream',
392 many=False,
393 default=1,
394 minimum=1,
395 makeStream=True,
396 streamOffset=0.75,
397 streamName='S',
398 ),
399 'retentate': JsonPortConfig(
400 displayName='Retentate',
401 type='outlet',
402 streamType='stream',
403 many=False,
404 default=1,
405 minimum=1,
406 makeStream=True,
407 streamOffset=0.75,
408 streamName='Retentate S',
409 ),
410 'permeate': JsonPortConfig(
411 displayName='Permeate',
412 type='outlet',
413 streamType='stream',
414 many=False,
415 default=1,
416 minimum=1,
417 makeStream=True,
418 streamOffset=0.75,
419 streamName='Permeate S',
420 ),
421 },
422 propertyPackagePorts={
423 '': ['inlet', 'retentate', 'permeate'],
424 },
425 graphicObject=JsonGraphicObjectConfig(
426 kind='unitop_graphic',
427 ),
428 indexSets=[],
429 properties={
430 'A_comp': JsonPropertyConfig(
431 propertySetGroup='default',
432 displayName='Membrane Water Permeability (A)',
433 indexSets=None,
434 sumToOne=False,
435 value=None,
436 unit=None,
437 unitType='solventPermeability',
438 description=None,
439 type='numeric',
440 many=False,
441 default=1,
442 options={},
443 hasTimeIndex=True,
444 ),
445 'B_comp': JsonPropertyConfig(
446 propertySetGroup='default',
447 displayName='Membrane Salt Permeability (B)',
448 indexSets=None,
449 sumToOne=False,
450 value=None,
451 unit=None,
452 unitType='velocity',
453 description=None,
454 type='numeric',
455 many=False,
456 default=1,
457 options={},
458 hasTimeIndex=True,
459 ),
460 'permeate.pressure': JsonPropertyConfig(
461 propertySetGroup='default',
462 displayName='Permeate Pressure',
463 indexSets=None,
464 sumToOne=False,
465 value=None,
466 unit=None,
467 unitType='pressure',
468 description=None,
469 type='numeric',
470 many=False,
471 default=1,
472 options={},
473 hasTimeIndex=True,
474 ),
475 'area': JsonPropertyConfig(
476 propertySetGroup='default',
477 displayName='Membrane Area',
478 indexSets=None,
479 sumToOne=False,
480 value=None,
481 unit=None,
482 unitType='area',
483 description=None,
484 type='numeric',
485 many=False,
486 default=1,
487 options={},
488 hasTimeIndex=True,
489 ),
490 'feed_side.spacer_porosity': JsonPropertyConfig(
491 propertySetGroup='default',
492 displayName='Feed Spacer Porosity',
493 indexSets=None,
494 sumToOne=False,
495 value=None,
496 unit=None,
497 unitType='ratio',
498 description=None,
499 type='numeric',
500 many=False,
501 default=1,
502 options={},
503 hasTimeIndex=True,
504 ),
505 'feed_side.channel_height': JsonPropertyConfig(
506 propertySetGroup='default',
507 displayName='Feed Channel Height',
508 indexSets=None,
509 sumToOne=False,
510 value=None,
511 unit=None,
512 unitType='distance',
513 description=None,
514 type='numeric',
515 many=False,
516 default=1,
517 options={},
518 hasTimeIndex=True,
519 ),
520 'length': JsonPropertyConfig(
521 propertySetGroup='default',
522 displayName='Membrane Length',
523 indexSets=None,
524 sumToOne=False,
525 value=None,
526 unit=None,
527 unitType='distance',
528 description=None,
529 type='numeric',
530 many=False,
531 default=1,
532 options={},
533 hasTimeIndex=True,
534 ),
535 'width': JsonPropertyConfig(
536 propertySetGroup='default',
537 displayName='Membrane Width',
538 indexSets=None,
539 sumToOne=False,
540 value=None,
541 unit=None,
542 unitType='distance',
543 description=None,
544 type='numeric',
545 many=False,
546 default=1,
547 options={},
548 hasTimeIndex=True,
549 ),
550 'inlet_reynolds_number': JsonPropertyConfig(
551 propertySetGroup='default',
552 displayName='Inlet Reynolds Number',
553 indexSets=None,
554 sumToOne=False,
555 value=None,
556 unit=None,
557 unitType='dimensionless',
558 description=None,
559 type='numeric',
560 many=False,
561 default=1,
562 options={},
563 hasTimeIndex=True,
564 ),
565 'volumetric_recovery': JsonPropertyConfig(
566 propertySetGroup='default',
567 displayName='Volumetric Recovery',
568 indexSets=None,
569 sumToOne=False,
570 value=None,
571 unit=None,
572 unitType='ratio',
573 description=None,
574 type='numeric',
575 many=False,
576 default=1,
577 options={},
578 hasTimeIndex=True,
579 ),
580 'tds_mass_recovery': JsonPropertyConfig(
581 propertySetGroup='default',
582 displayName='Mass Recovery',
583 indexSets=None,
584 sumToOne=False,
585 value=None,
586 unit=None,
587 unitType='ratio',
588 description=None,
589 type='numeric',
590 many=False,
591 default=1,
592 options={},
593 hasTimeIndex=True,
594 ),
595 'tds_observed_rejection': JsonPropertyConfig(
596 propertySetGroup='default',
597 displayName='Observed Solute Rejection',
598 indexSets=None,
599 sumToOne=False,
600 value=None,
601 unit=None,
602 unitType='ratio',
603 description=None,
604 type='numeric',
605 many=False,
606 default=1,
607 options={},
608 hasTimeIndex=True,
609 ),
610 'average_water_mass_flux': JsonPropertyConfig(
611 propertySetGroup='default',
612 displayName='Average Mass Flux',
613 indexSets=None,
614 sumToOne=False,
615 value=None,
616 unit=None,
617 unitType='massFlux',
618 description=None,
619 type='numeric',
620 many=False,
621 default=1,
622 options={},
623 hasTimeIndex=True,
624 ),
625 'over_pressure_ratio': JsonPropertyConfig(
626 propertySetGroup='default',
627 displayName='Over Pressure Ratio',
628 indexSets=None,
629 sumToOne=False,
630 value=None,
631 unit=None,
632 unitType='ratio',
633 description=None,
634 type='numeric',
635 many=False,
636 default=1,
637 options={},
638 hasTimeIndex=True,
639 ),
640 },
641 propertySetGroups={
642 'default': JsonPropertySetGroupConfig(
643 type='stateVars',
644 displayName='Properties',
645 stateVars=[
646 'A_comp',
647 'B_comp',
648 'permeate.pressure',
649 'area',
650 'feed_side.spacer_porosity',
651 'feed_side.channel_height',
652 'length',
653 ],
654 toggle=None,
655 ),
656 },
657 keyProperties=[
658 'volumetric_recovery',
659 'tds_mass_recovery',
660 'tds_observed_rejection',
661 'average_water_mass_flux',
662 'over_pressure_ratio',
663 ],
664 splitterFractionName=None,
665 idaesAdapter=JsonIdaesAdapterConfig(
666 constructor='ahuora_builder.custom.watertap.reverse_osmosis_0d.ReverseOsmosis0D',
667 args={
668 'property_package': JsonAdapterArgConfig(
669 kind='property_package',
670 label=None,
671 ),
672 },
673 ports=None,
674 properties=None,
675 ),
676 frontend=JsonFrontendConfig(
677 showInPanel=True,
678 variant=None,
679 ),
680 )