Mask the loss weight alongside the error (use_f_weights is broken)

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):
     """

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)