diff --git a/cppcore/src/system/HoppingBlocks.cpp b/cppcore/src/system/HoppingBlocks.cpp index 4f0c4d27..c2041a7d 100644 --- a/cppcore/src/system/HoppingBlocks.cpp +++ b/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()]; block.reserve(block.size() + rows.size()); diff --git a/pybinding/modifier.py b/pybinding/modifier.py index f458dd99..ac25ab36 100644 --- a/pybinding/modifier.py +++ b/pybinding/modifier.py @@ -418,7 +418,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 @@ -429,6 +429,8 @@ 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) @@ -436,7 +438,8 @@ def _make_generator(func, kind, name, energy, keywords): 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) @@ -519,7 +522,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") diff --git a/tests/test_modifiers.py b/tests/test_modifiers.py index bd3fc3dc..d0cf5346 100644 --- a/tests/test_modifiers.py +++ b/tests/test_modifiers.py @@ -334,6 +334,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"""