Allow returning bool arrays in hopping generators

cppcore/src/system/HoppingBlocks.cpp

@@ -57,7 +57,11 @@ void HoppingBlocks::reserve(ArrayXi const& counts) {
 }
 
 void HoppingBlocks::append(HopID family_id, ArrayXi&& rows, ArrayXi&& cols) {
-    assert(rows.size() == cols.size());
+    if (rows.size() != cols.size()) {
+        throw std::runtime_error("When generating hoppings, the number of "
+                                 "`from` and `to` indices must be equal");
+    }
+
     auto& block = blocks[family_id.as<size_t>()];
     block.reserve(block.size() + rows.size());
 

pybinding/modifier.py

@@ -416,7 +416,7 @@ def f(energy):
     return f
 
 
-def _make_generator(func, kind, name, energy, keywords):
+def _make_generator(func, kind, name, energy, keywords, process_result=lambda x, *_: x):
     """Turn a regular function into a generator of the desired kind
 
     Parameters
@@ -427,14 +427,17 @@ def _make_generator(func, kind, name, energy, keywords):
         Modifier base class.
     keywords : str
         String of comma separated names: the expected arguments of a modifier function.
+    process_result : Callable
+        Apply additional processing on the generator result
     """
     keywords = [word.strip() for word in keywords.split(",")]
     _check_modifier_spec(func, keywords)
     requested_argnames = tuple(inspect.signature(func).parameters.keys())
 
     def generator_func(*args):
         requested_kwargs = _process_modifier_args(args, keywords, requested_argnames)
-        return func(**requested_kwargs)
+        result = func(**requested_kwargs)
+        return process_result(result, *args)
 
     class Generator(kind):
         callsig = getattr(func, 'callsig', None)
@@ -517,7 +520,16 @@ def hopping_generator(name, energy):
     The function must return:
 
     Tuple[np.ndarray, np.ndarray]
-        Arrays of index pairs which form the new hoppings.
+        A pair of arrays of indices which form the new hoppings.
     """
-    return functools.partial(_make_generator, kind=_cpp.HoppingGenerator,
-                             name=name, energy=energy, keywords="system, x, y, z")
+    def process_result(result, system):
+        def process(v):
+            v = np.asarray(v)
+            if v.dtype == np.bool_ and v.size == system.num_sites:
+                return np.flatnonzero(v)
+            else:
+                return v
+        return tuple(process(v) for v in result)
+
+    return functools.partial(_make_generator, kind=_cpp.HoppingGenerator, name=name, energy=energy,
+                             process_result=process_result, keywords="system, x, y, z")

tests/test_modifiers.py

@@ -335,6 +335,60 @@ def onsite_offset(energy):
     expected = pb.Model(graphene.monolayer(2), graphene.hexagon_ac(1))
     assert pytest.fuzzy_equal(model.hamiltonian, expected.hamiltonian)
 
+    @pb.hopping_generator("t_new", energy=1.0)
+    def bad_generator():
+        """Different array lengths"""
+        return [0, 1, 2], [0, 1]
+
+    model = pb.Model(graphene.monolayer(), pb.primitive(3, 3), bad_generator)
+    with pytest.raises(RuntimeError) as excinfo:
+        model.eval()
+    assert "the number of `from` and `to` indices must be equal" in str(excinfo.value)
+
+
+def test_site_and_hopping_interaction():
+    """Add a new row of sites and connect them just like the rest of the lattice"""
+    d = 1.0
+    v = 1.5
+    t = 1.0
+
+    def square_lattice():
+        lat = pb.Lattice(a1=[d, 0], a2=[0, d])
+        lat.add_sublattices(("A", [0, 0], v))
+        lat.add_hoppings(([0, 1], "A", "A", t),
+                         ([1, 0], "A", "A", t))
+        return lat
+
+    @pb.site_generator(name="B", energy=v)
+    def edge_sites(system):
+        edge_atoms = system.count_neighbors() < 4
+        x, y, z = (v[edge_atoms] for v in system.positions)
+
+        top_edge_only = np.isclose(y, y.max())
+        x, y, z = (v[top_edge_only] for v in (x, y, z))
+
+        y += d
+        return x, y, z
+
+    @pb.hopping_generator(name="t_edge", energy=t)
+    def edge_hoppings(system, y):
+        new_sites = system.sub == "B"
+        edge_atoms = np.logical_and(system.sub != "B", system.count_neighbors() < 4)
+        top_edge_only = np.logical_and(edge_atoms, np.isclose(y, y[edge_atoms].max()))
+        return new_sites, top_edge_only
+
+    @pb.hopping_generator(name="t_edge2", energy=t)
+    def edge_hoppings2(system):
+        edge_idx = np.flatnonzero(system.sub == "B")
+        to_idx = edge_idx[1:]
+        from_idx = edge_idx[:-1]
+        return to_idx, from_idx
+
+    model = pb.Model(square_lattice(), pb.primitive(6, 4),
+                     edge_sites, edge_hoppings, edge_hoppings2)
+    expected = pb.Model(square_lattice(), pb.primitive(6, 5))
+    assert pytest.fuzzy_equal(model.hamiltonian, expected.hamiltonian)
+
 
 def test_wrapper_return():
     """Make sure the wrapper return type conversion is working"""