FPIDT: federated physics-informed digital twins for reverse-osmosis desalination

Models from the paper Federated Physics-Informed Digital Twins for Privacy-Preserving Performance Prediction of Reverse Osmosis Desalination (Emad Eldin, Baroud, Atitallah, AL-Ghamdi). They predict water flux and specific energy consumption (SEC) of RO membranes from four operating variables, and were trained with federated learning across nine real pilot-plant clients, so no client's raw data left that client. All training data are public experimental records from two U.S. national-laboratory RO pilots; nothing is simulated.

Headline result. With a client's full history, a network trained on that client alone is the most accurate. Personalisation brings federated models close to it and well ahead of one centralized model for all plants. The physics loss improves the shared model, keeps its responses physically correct, and helps where local data are missing (a plant left out of training, or one with only minutes of history). Use the per-client models for known plants and global_pi_fedavg for a new one.

1. The problem

RO plants log pressure, flow, conductivity and power continuously. A model that predicts how much water a membrane will produce, and at what energy cost, supports operation and what-if analysis (a digital twin). Training it on one plant limits it to that plant's operating range; pooling data across plants is often impossible because operators cannot share raw records. Federated learning trains one model across sites while raw data stay local, but plants differ in membrane, salinity and control strategy (the data are non-IID). The proposal tested here adds two remedies: a physics-informed loss that holds every client to the solution–diffusion law of RO, and personalisation (local fine-tuning of the shared model).

2. Data

From 944 raw files (6,172,941 logged rows at 2–4 Hz), filtration-phase rows with feed pressure above 3.5 bar and flux above 0.5 LMH were kept and averaged into 10-second windows (216,909 windows). Each client is split in time: first 70% train, next 10% validation, last 20% test, with 5-minute gaps between blocks, so every test is a prediction of later operation.

3. Features and targets

Leakage guard. Never used as inputs: permeate flow (flux × area), NAWI's make-up feed flow (equal to permeate flow during closed-circuit filtration), cumulative permeate volume, and — for the energy model — the pump powers that define SEC.

4. Exploratory analysis

The clients are strongly non-IID. NAWI clients run at a median of 17.6 bar on brackish loops of 12.1 mS/cm; NREL clients at 54.6 bar on seawater of 56.8 mS/cm.

Control strategy shapes the data. In the extremum-seeking, predictive-control and Bayesian-optimisation campaigns the controller holds flux at a set point while pressure rises with the loop's salinity; in the open-loop campaigns and at NREL, flux follows the applied pressure.

The physics holds on both pilots. Each client's solution–diffusion fit J = A (P − kC) gives a water permeability A of 1.34–2.13 LMH/bar on the brackish NAWI membrane and 0.79–0.98 LMH/bar on the seawater NREL membrane, and an osmotic coefficient k of 0.40–0.55 bar per mS/cm on every client — physically plausible for NaCl.

Energy. SEC falls with flux on every NAWI client, as fixed pump overheads are spread over more water.

5. Models in this repository

All models share one architecture: an MLP [4, 64, 64, 1] with SiLU activations (4,545 parameters for flux, 4,609 for energy). Per-client files stack one model per client in the order of config.json → clients. Each exported model was retrained with seed 0 and re-scored: every per-client test RMSE matches the paper's seed-0 run exactly (reproduction_check in config.json).

Which to use. For a plant that is one of the clients, personalised_fpidt or local_nn (the most accurate; see results). For a new plant with no data, global_pi_fedavg (the physics-informed shared model, best of the privacy-preserving options on unseen clients for flux).

6. Training

• Federated averaging: 60 rounds; each round every client starts from the global model and takes 100 Adam steps (batch 256) on its own data; the server averages parameters weighted by training size; the round with the lowest mean per-client validation RMSE is kept.
• Physics loss (PI variants, weight λ): squared deviation from the client's own solution–diffusion fit, plus hinge penalties when +0.5 bar lowers the predicted flux or +0.5 mS/cm raises it. For energy: deviation from a local grey-box fit SEC = aP + bQ/J + c plus a hinge on predictions below the hydraulic bound SEC ≥ P/36 (P in bar). Each client's physics parameters stay local.
• Personalisation: fine-tune the global model on the client's own data for up to 300 steps (lr 5e-4), keeping the state with the lowest validation error on that client.
• Standardisation: global means and standard deviations from per-client sums (stored in config.json).
• Hyper-parameters, chosen on validation blocks only (grids widened until the optimum was interior): flux λ = 1, FedProx μ = 1, fine-tuning λ = 0; energy λ = 0.1, μ = 1, fine-tuning λ = 0.01.
• Each message carries 4,545 parameters (~17.8 kB); a nine-client round moves 0.31 MB in both directions.

7. Results (paper; 5 seeds, test blocks)

Statistics treat clients as the independent units: per-client RMSE is averaged over seeds, and two methods are compared with an exact sign-flip permutation test over clients, Holm-adjusted. With nine flux clients the smallest attainable p is 0.004; with six energy clients it is 0.031.

Water flux (9 clients)

Specific energy (6 NAWI clients)

"Physics violations": share of test points where +0.5 bar lowers, or +0.5 mS/cm raises, the predicted flux (for energy: predictions below P/36).

What each ingredient does (2×2 design)

Flux

Energy

A plant that contributed no data (leave one client out; 2 seeds, 40 rounds)

Flux (LMH)

PI-FedAvg − FedAvg: -2.12 (sign-flip p = 0.04).

Energy (kWh/m³)

PI-FedAvg − FedAvg: +0.01 (sign-flip p = 0.19).

Plants with little history (flux; 3 seeds)

Every client keeps only its most recent training windows; validation and test blocks are unchanged.

With 5 minutes of history FPIDT has the lowest mean error, but beats the local model on only 5/9 clients (Holm p = 1.00): it helps clients whose short history does not cover their later operating range and hurts set-point campaigns whose flux is fixed by the controller.

8. How to use

from huggingface_hub import snapshot_download
import sys
path = snapshot_download("FatimahEmadEldin/fpidt-ro-desalination", token=True)
sys.path.insert(0, path)
from fpidt_model import load

# personalised model for a known client: inputs P_bar, C_bulk, T_C, Q_cross (raw units)
m = load("flux", "personalised_fpidt", client="N4_openloop_aug")
print(m.predict([[18.2, 17.1, 22.6, 49.9]]))          # water flux, LMH

# shared physics-informed model for a new plant
g = load("flux", "global_pi_fedavg")
print(g.predict([[50.0, 57.0, 20.0, 20.4]]))

# energy: inputs P_bar, C_bulk, T_C, Q_cross, J_LMH
e = load("energy", "personalised_fpidt", client="N5_mpc")
print(e.predict([[22.9, 16.3, 21.7, 45.3, 24.6]]))    # SEC, kWh/m^3

Requires torch, safetensors, numpy. config.json in each task folder lists feature order, units, standardisation, clients, physics fits, training settings and per-client test RMSE for every model.

9. Limitations and intended use

• Research models, trained on two pilot rigs (six campaigns on one brackish rig, three regimes on one seawater rig), not on full-scale plants; they are not validated for operational control.
• Inputs outside the training range (pressure, salinity, crossflow) are extrapolation; the global physics-informed model degrades more gracefully than the local ones, but none is guaranteed.
• On set-point campaigns (N3, N5, N6) flux is held by the controller, so flux predictions there mostly reproduce the set point; locally trained models can learn the controller's confounded sign.
• Energy exists only for the six NAWI clients; NREL clients have flux models only.
• Privacy is architectural (raw data stay local); no secure aggregation or differential privacy was used, and model parameters can leak information about training data.
• Membrane fouling is not modelled: the public data carry no fouling labels.
• License: other — to be set by the authors. Training data are subject to the terms of the NAWI DAMS and MHKDR repositories.

10. Citation

Paper: F. Emad Eldin, B. Baroud, N. Atitallah, A. S. AL-Malaise AL-Ghamdi, Federated Physics-Informed Digital Twins for Privacy-Preserving Performance Prediction of Reverse Osmosis Desalination (2026).

Data: K. Villez et al., NAWI Seedling Project: Foundational Control Methods for Water Treatment Systems — Online Actuator and Sensor Data, NAWI Water DAMS / ORNL, doi:10.7481/1963162. National Renewable Energy Laboratory, Experimental Wave Desalination Data with Variable Feed and Pressure Input, MHKDR submission 649 (2025); see also K. A. Sitterley, Z. Binger, D. S. Jenne, Membranes 15(8), 243 (2025), doi:10.3390/membranes15080243.