From d94cdd9735dbbab853920d710c6150aee3c00d74 Mon Sep 17 00:00:00 2001 From: Roy Kid Date: Wed, 16 Sep 2026 15:46:58 +0200 Subject: [PATCH] Mask the loss weight alongside the error `tf.boolean_mask` flattens what it selects, so a force error of shape [n_atoms, 3] comes back as [n_atoms*3]. The weight did not travel through the same selection, so `error**2 * weight` tried to broadcast [n_atoms*3] against [n_atoms, 3] and the graph failed to build: Incompatible shapes: [579] vs. [193,3] The force error is always masked -- `f_mask` is an all-true fill -- so this hit every run with `use_f_weights: True`, i.e. the option has never worked since boolean masking arrived in 1586487 (2020). Scalar weights broadcast fine, which is why the energy path (mask None unless `max_energy` is set) never showed it. Select the weight with the same mask when it is a tensor of the error's rank, and add a regression test that trains and evaluates with per-component force weights. --- pinn/models/base.py | 22 ++++++++++++++++++++-- tests/test_potential.py | 41 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+), 2 deletions(-) diff --git a/pinn/models/base.py b/pinn/models/base.py index 4bd4a88b..2ddbb1b3 100644 --- a/pinn/models/base.py +++ b/pinn/models/base.py @@ -142,7 +142,7 @@ def add_error(self, tag, data, pred, mask=None, weight=1.0, tf.compat.v1.summary.scalar(f'{tag}_MAE', mae) tf.compat.v1.summary.scalar(f'{tag}_RMSE', rmse) if mask is not None: - error = tf.boolean_mask(error, mask) + error, weight = _apply_mask(error, weight, mask) if use_error: loss = tf.reduce_mean(error**2 * weight) tf.compat.v1.summary.scalar(f'{tag}_LOSS', loss) @@ -153,13 +153,31 @@ def add_error(self, tag, data, pred, mask=None, weight=1.0, self.METRICS[f'METRICS/{tag}_MAE'] = tf.compat.v1.metrics.mean_absolute_error(data, pred) self.METRICS[f'METRICS/{tag}_RMSE'] = tf.compat.v1.metrics.root_mean_squared_error(data, pred) if mask is not None: - error = tf.boolean_mask(error, mask) + error, weight = _apply_mask(error, weight, mask) if use_error: loss = tf.reduce_mean(error**2 * weight) self.METRICS[f'METRICS/{tag}_LOSS'] = tf.compat.v1.metrics.mean(loss) self.LOSS.append(loss) +def _apply_mask(error, weight, mask): + """Select the kept entries of an error and of a per-entry weight. + + ``tf.boolean_mask`` flattens what it selects, so a force error of shape + ``[n_atoms, 3]`` comes back as ``[n_kept]``. A weight of the same rank -- + ``f_weights`` under ``use_f_weights`` -- has to travel through the same + selection, or the ``error**2 * weight`` that follows tries to broadcast + ``[n_kept]`` against ``[n_atoms, 3]`` and the graph fails to build. Scalar + weights (a plain loss multiplier) broadcast fine and are passed through. + + Returns: + tuple: the masked error and the weight to multiply it by. + """ + if tf.is_tensor(weight) and weight.shape.rank == error.shape.rank: + weight = tf.boolean_mask(weight, mask) + return tf.boolean_mask(error, mask), weight + + @pi_named('TRAIN_OP') def get_train_op(optimizer, metrics, tvars, separate_errors=False): """ diff --git a/tests/test_potential.py b/tests/test_potential.py index b1fe6bdf..9e512223 100644 --- a/tests/test_potential.py +++ b/tests/test_potential.py @@ -262,3 +262,44 @@ def test(): return load_numpy(data).apply(sparse_batch(10)) de = e_pred[-1] - e_pred[0] int_p = np.trapz(p_pred, x=l_range**3) assert np.allclose(de, int_p, rtol=1e-2) + + +@pytest.mark.forked +def test_use_f_weights_trains(): + """Per-component force weights survive the loss mask. + + The force error is always masked (`f_mask` is an all-true fill), and + `tf.boolean_mask` flattens it; a `[n_atoms, 3]` weight that did not travel + through the same selection used to make the graph fail to build with + `Incompatible shapes: [n_atoms*3] vs [n_atoms, 3]`. + """ + testpath = tempfile.mkdtemp() + n, na = 8, 4 + rng = np.random.default_rng(0) + data = { + 'coord': rng.normal(0, 2.0, (n, na, 3)).astype(np.float32), + 'elems': np.ones((n, na), dtype=np.int32), + 'e_data': rng.normal(-20, 1, n).astype(np.float32), + 'f_data': rng.normal(0, 0.5, (n, na, 3)).astype(np.float32), + 'f_weights': rng.uniform(0.5, 1.5, (n, na, 3)).astype(np.float32), + } + params = { + 'model_dir': testpath, + 'network': { + 'name': 'PiNet', + 'params': {'ii_nodes': [4, 4], 'pi_nodes': [4, 4], 'pp_nodes': [4, 4], + 'out_nodes': [4], 'depth': 2, 'rc': 4., 'n_basis': 4, + 'atom_types': [1]}}, + 'model': { + 'name': 'potential_model', + 'params': {'use_force': True, 'use_f_weights': True, + 'f_loss_multiplier': 100.0}}} + + import pinn + model = pinn.get_model(params) + model.train( + input_fn=lambda: load_numpy(data, splits=None).apply(sparse_batch(2)).repeat(), + max_steps=2) + model.evaluate( + input_fn=lambda: load_numpy(data, splits=None).apply(sparse_batch(2))) + rmtree(testpath)