Add site generator

cppcore/include/Model.hpp

@@ -38,6 +38,7 @@ class Model {
     void add(OnsiteModifier const& m);
     void add(HoppingModifier const& m);
 
+    void add(SiteGenerator const& g);
     void add(HoppingGenerator const& g);
 
     void set_wave_vector(Cartesian const& k);

cppcore/include/system/CompressedSublattices.hpp

@@ -70,7 +70,7 @@ class CompressedSublattices {
                           ArrayXi const& orbital_counts);
 
     /// Start a new sublattice block or increment the site count for the existing block
-    void add(SiteID id, idx_t norb);
+    void add(SiteID id, idx_t norb, idx_t count = 1);
 
     /// Remove sites for which `keep == false`
     void filter(VectorX<bool> const& keep);

cppcore/include/system/HoppingBlocks.hpp

@@ -6,6 +6,9 @@
 
 namespace cpb {
 
+/// Alternative CSR mapping of hopping IDs
+using HoppingCSR = SparseMatrixX<storage_idx_t>;
+
 /**
  A simple row and column index pair
  */
@@ -124,8 +127,11 @@ class HoppingBlocks {
     /// Remove sites for which `keep == false`
     void filter(VectorX<bool> const& keep);
 
+    /// Account for the addition of new sites (no new hoppings)
+    void add_sites(idx_t num_new_sites);
+
     /// Return the matrix in the CSR sparse matrix format
-    SparseMatrixX<storage_idx_t> tocsr() const;
+    HoppingCSR tocsr() const;
 
 private:
     idx_t num_sites; ///< number of lattice sites, i.e. the size of the square matrix

cppcore/include/system/Registry.hpp

@@ -64,6 +64,7 @@ void check_onsite_energy(MatrixXcd const& energy);
 /// Convert the onsite energy into the canonical format
 MatrixXcd canonical_onsite_energy(std::complex<double> energy);
 MatrixXcd canonical_onsite_energy(VectorXd const& energy);
+inline MatrixXcd canonical_onsite_energy(MatrixXcd const& energy) { return energy; }
 
 /// Check that the hopping energy matrix satisfies all the requirements
 void check_hopping_energy(MatrixXcd const& energy);

cppcore/include/system/StructureModifiers.hpp

@@ -1,5 +1,6 @@
 #pragma once
-#include "Lattice.hpp"
+#include "CompressedSublattices.hpp"
+#include "HoppingBlocks.hpp"
 
 #include "numeric/dense.hpp"
 
@@ -45,6 +46,29 @@ struct SubIdRef {
     std::unordered_map<std::string, storage_idx_t> name_map;
 };
 
+/**
+ Introduces a new site family (with new sub_id)
+
+ This can be used to create new sites independent of the translations of the main unit cell
+ as define by the `Lattice` class. It's useful for disorder or terminating system edges with
+ atoms of a different element.
+ */
+class SiteGenerator {
+public:
+    using Function = std::function<
+        CartesianArray(CartesianArrayConstRef, CompressedSublattices const&, HoppingBlocks const&)
+    >;
+
+    std::string name; ///< friendly site family identifier
+    MatrixXcd energy; ///< onsite energy - also added to the site registry
+    Function make; ///< function which will generate the new site positions
+
+    SiteGenerator(string_view name, MatrixXcd const& energy, Function const& make)
+        : name(name), energy(energy), make(make) {}
+
+    explicit operator bool() const { return static_cast<bool>(make); }
+};
+
 /**
  Introduces a new hopping family (with new hop_id) via a list of index pairs
 
@@ -77,12 +101,15 @@ void apply(PositionModifier const& m, Foundation& f);
 template<class M> void apply(M const&, System&) {}
 void apply(SiteStateModifier const& m, System& s);
 void apply(PositionModifier const& m, System& s);
+void apply(SiteGenerator const& g, System& s);
 void apply(HoppingGenerator const& g, System& s);
 
 template<class M> constexpr bool requires_system(M const&) { return false; }
+constexpr bool requires_system(SiteGenerator const&) { return true; }
 constexpr bool requires_system(HoppingGenerator const&) { return true; }
 
 template<class M> constexpr bool is_generator(M const&) { return false; }
+constexpr bool is_generator(SiteGenerator const&) { return true; }
 constexpr bool is_generator(HoppingGenerator const&) { return true; }
 
 /**

cppcore/src/Model.cpp

@@ -72,6 +72,12 @@ void Model::add(HoppingModifier const& m) {
     clear_hamiltonian();
 }
 
+void Model::add(SiteGenerator const& g) {
+    structure_modifiers.emplace_back(g);
+    site_registry.register_family(g.name, g.energy);
+    clear_structure();
+}
+
 void Model::add(HoppingGenerator const& g) {
     structure_modifiers.emplace_back(g);
     hopping_registry.register_family(g.name, MatrixXcd::Constant(1, 1, g.energy));

cppcore/src/system/CompressedSublattices.cpp

@@ -12,11 +12,11 @@ CompressedSublattices::CompressedSublattices(ArrayXi const& alias_ids, ArrayXi c
     }
 }
 
-void CompressedSublattices::add(SiteID id, idx_t norb) {
+void CompressedSublattices::add(SiteID id, idx_t norb, idx_t count) {
     if (data.empty() || data.back().id != id) {
-        data.push_back({id, 1, static_cast<storage_idx_t>(norb)});
+        data.push_back({id, static_cast<storage_idx_t>(count), static_cast<storage_idx_t>(norb)});
     } else {
-        data.back().num_sites += 1;
+        data.back().num_sites += static_cast<storage_idx_t>(count);
     }
 }
 

cppcore/src/system/HoppingBlocks.cpp

@@ -73,8 +73,12 @@ void HoppingBlocks::filter(VectorX<bool> const& keep) {
     }
 }
 
-SparseMatrixX<storage_idx_t> HoppingBlocks::tocsr() const {
-    auto csr = SparseMatrixX<storage_idx_t>(num_sites, num_sites);
+void HoppingBlocks::add_sites(idx_t num_new_sites) {
+    num_sites += num_new_sites;
+}
+
+HoppingCSR HoppingBlocks::tocsr() const {
+    auto csr = HoppingCSR(num_sites, num_sites);
     csr.reserve(nnz());
     for (auto const& block : *this) {
         for (auto const& coo : block.coordinates()) {

cppcore/src/system/StructureModifiers.cpp

@@ -47,6 +47,17 @@ void apply(PositionModifier const& m, System& s) {
     }
 }
 
+void apply(SiteGenerator const& g, System& s) {
+    detail::remove_invalid(s);
+
+    auto const new_positions = g.make(s.positions, s.compressed_sublattices, s.hopping_blocks);
+    auto const norb = g.energy.rows();
+    auto const nsites = new_positions.size();
+    s.compressed_sublattices.add(s.site_registry.id(g.name), norb, nsites);
+    s.hopping_blocks.add_sites(nsites);
+    s.positions = concat(s.positions, new_positions);
+}
+
 void apply(HoppingGenerator const& g, System& s) {
     detail::remove_invalid(s);
 

cppcore/tests/test_modifiers.cpp

@@ -177,6 +177,69 @@ TEST_CASE("HoppingEnergyModifier") {
     REQUIRE(h.non_zeros() == 0);
 }
 
+TEST_CASE("SiteGenerator") {
+    auto model = Model([]{
+        auto lattice = Lattice({1, 0, 0}, {0, 1, 0});
+        lattice.add_sublattice("A", {0, 0, 0});
+        lattice.add_sublattice("B", {0, 0, 0});
+        lattice.register_hopping_energy("t1", 1.0);
+        return lattice;
+    }());
+    REQUIRE_FALSE(model.is_complex());
+    REQUIRE(model.get_lattice().get_hoppings().size() == 1);
+    REQUIRE(model.system()->hopping_blocks.nnz() == 0);
+
+    SECTION("Errors") {
+        auto const noop = [](CartesianArrayConstRef, CompressedSublattices const&,
+                             HoppingBlocks const&) { return CartesianArray(); };
+
+        auto const complex_vector = MatrixXcd::Constant(1, 2, 2.0);
+        REQUIRE_THROWS_WITH(model.add(SiteGenerator("C", complex_vector, noop)),
+                            Catch::Contains("must be a real vector or a square matrix"));
+
+        auto const complex_matrix = MatrixXcd::Constant(2, 2, {1.0, 1.0});
+        REQUIRE_THROWS_WITH(model.add(SiteGenerator("C", complex_matrix, noop)),
+                            Catch::Contains("diagonal of the onsite hopping term must be real"));
+    }
+
+    SECTION("Structure") {
+        auto const energy = MatrixXcd::Constant(1, 1, 2.0);
+        model.add(SiteGenerator("C", energy, [](CartesianArrayConstRef,
+                                                CompressedSublattices const&,
+                                                HoppingBlocks const&) {
+            auto const size = 5;
+            auto x = ArrayXf::Constant(size, 1);
+            auto y = ArrayXf::LinSpaced(size, 1, 5);
+            auto z = ArrayXf::Constant(size, 0);
+            return CartesianArray(x, y, z);
+        }));
+
+        REQUIRE_FALSE(model.is_complex());
+        REQUIRE(model.get_lattice().get_sublattices().size() == 2);
+        REQUIRE(model.get_site_registry().size() == 3);
+        REQUIRE(model.system()->compressed_sublattices.alias_ids().size() == 3);
+        REQUIRE(model.system()->num_sites() == 7);
+
+        REQUIRE(model.system()->positions[0].isApprox(Cartesian{0, 0, 0}));
+        REQUIRE(model.system()->positions[1].isApprox(Cartesian{0, 0, 0}));
+        REQUIRE(model.system()->positions[2].isApprox(Cartesian{1, 1, 0}));
+        REQUIRE(model.system()->positions[3].isApprox(Cartesian{1, 2, 0}));
+        REQUIRE(model.system()->positions[4].isApprox(Cartesian{1, 3, 0}));
+        REQUIRE(model.system()->positions[5].isApprox(Cartesian{1, 4, 0}));
+        REQUIRE(model.system()->positions[6].isApprox(Cartesian{1, 5, 0}));
+
+        auto const names = model.get_site_registry().name_map();
+        auto const it = names.find("C");
+        REQUIRE(it != names.end());
+
+        auto const id = it->second;
+        auto const& cs = model.system()->compressed_sublattices;
+        REQUIRE(cs.alias_ids()[2] == id);
+        REQUIRE(cs.orbital_counts()[2] == 1);
+        REQUIRE(cs.site_counts()[2] == 5);
+    }
+}
+
 TEST_CASE("HoppingGenerator") {
     auto model = Model([]{
         auto lattice = Lattice({1, 0, 0}, {0, 1, 0});

cppmodule/src/model.cpp

@@ -12,6 +12,7 @@ void wrap_model(py::module& m) {
         .def("add", &Model::add | resolve<PositionModifier const&>())
         .def("add", &Model::add | resolve<OnsiteModifier const&>())
         .def("add", &Model::add | resolve<HoppingModifier const&>())
+        .def("add", &Model::add | resolve<SiteGenerator const&>())
         .def("add", &Model::add | resolve<HoppingGenerator const&>())
         .def("attach_lead", &Model::attach_lead)
         .def("set_wave_vector", &Model::set_wave_vector, "k"_a, R"(
@@ -22,6 +23,7 @@ void wrap_model(py::module& m) {
             k : array_like
                 Wave vector in reciprocal space.
         )")
+        .def_property_readonly("lattice", &Model::get_lattice)
         .def_property_readonly("system", &Model::system)
         .def_property_readonly("raw_hamiltonian", &Model::hamiltonian)
         .def_property_readonly("hamiltonian", [](Model const& self) {

cppmodule/src/modifiers.cpp

@@ -39,6 +39,21 @@ void extract_modifier_result(T& v, py::object const& o) {
 
 } // anonymous namespace
 
+template<class T>
+void init_site_generator(SiteGenerator& self, string_view name, T const& energy, py::object make) {
+    new (&self) SiteGenerator(
+        name, detail::canonical_onsite_energy(energy),
+        [make](CartesianArrayConstRef p, CompressedSublattices const& s, HoppingBlocks const& h) {
+            py::gil_scoped_acquire guard{};
+            auto result = make(p.x(), p.y(), p.z(), &s, &h);
+            auto t = py::reinterpret_borrow<py::tuple>(result);
+            return CartesianArray(t[0].cast<ArrayXf>(),
+                                  t[1].cast<ArrayXf>(),
+                                  t[2].cast<ArrayXf>());
+        }
+    );
+};
+
 void wrap_modifiers(py::module& m) {
     py::class_<SubIdRef>(m, "SubIdRef")
         .def_property_readonly("ids", [](SubIdRef const& s) { return arrayref(s.ids); })
@@ -69,6 +84,11 @@ void wrap_modifiers(py::module& m) {
             });
         });
 
+    py::class_<SiteGenerator>(m, "SiteGenerator")
+        .def("__init__", init_site_generator<std::complex<double>>)
+        .def("__init__", init_site_generator<VectorXd>)
+        .def("__init__", init_site_generator<MatrixXcd>);
+
     py::class_<HoppingGenerator>(m, "HoppingGenerator")
         .def("__init__", [](HoppingGenerator& self, std::string const& name,
                             std::complex<double> energy, py::object make) {

cppmodule/src/system.cpp

@@ -47,7 +47,7 @@ void wrap_system(py::module& m) {
         });
 
     py::class_<HoppingBlocks>(m, "HoppingBlocks")
-        .def_property_readonly("nnz", [](HoppingBlocks const& hb) { return hb.nnz(); })
+        .def_property_readonly("nnz", &HoppingBlocks::nnz)
         .def("tocsr", [](HoppingBlocks const& hb) {
             auto type = py::module::import("pybinding.support.alias").attr("AliasCSRMatrix");
             return type(hb.tocsr(), "mapping"_a=hb.get_name_map());

pybinding/leads.py

@@ -132,7 +132,7 @@ def plot_contact(self, line_width=1.6, arrow_length=0.5,
             Color of the shaded area.
         """
         lead_spec = self.impl.spec
-        vectors = self.impl.system.lattice.vectors
+        vectors = self.lattice.vectors
         if len(lead_spec.shape.vertices) != 2 or len(vectors) != 2:
             raise RuntimeError("This only works for 2D systems")
 

pybinding/modifier.py

@@ -18,7 +18,7 @@
 
 __all__ = ['constant_potential', 'force_double_precision', 'force_complex_numbers',
            'hopping_energy_modifier', 'hopping_generator', 'onsite_energy_modifier',
-           'site_position_modifier', 'site_state_modifier']
+           'site_generator', 'site_position_modifier', 'site_state_modifier']
 
 
 def _process_modifier_args(args, keywords, requested_argnames):
@@ -455,6 +455,41 @@ def __call__(self, *args, **kwargs):
     return Generator()
 
 
+@decorator_decorator
+def site_generator(name, energy):
+    """Introduce a new site family (with a new `sub_id`) via a list of site positions
+
+    This can be used to create new sites independent of the main :class:`Lattice` definition.
+    It's especially useful for creating disorder or terminating system edges with atoms of
+    a different element.
+
+    Parameters
+    ----------
+    name : string
+        Friendly identifier for the new site family.
+    energy : Union[float, complex, array_like]
+        Base hopping energy value -- scalar or matrix.
+
+    Notes
+    -----
+    The function parameters must be a combination of any number of the following:
+
+    x, y, z : np.ndarray
+        Lattice site position.
+    sublattices : CompressedSublattices
+        TBD
+    hoppings : Hoppings
+        TBD
+
+    The function must return:
+
+    Tuple[np.ndarray, np.ndarray, np.ndarray]
+        Tuple of (x, y, z) arrays which indicate the positions of the new sites.
+    """
+    return functools.partial(_make_generator, kind=_cpp.SiteGenerator,
+                             name=name, energy=energy, keywords="x, y, z, sublattices, hoppings")
+
+
 @decorator_decorator
 def hopping_generator(name, energy):
     """Introduce a new hopping family (with a new `hop_id`) via a list of index pairs

tests/test_modifiers.py

@@ -295,6 +295,23 @@ def make_complex_double(energy):
     assert model.hamiltonian.dtype == np.complex128
 
 
+def test_site_generator():
+    """Generated some disordered sites"""
+    @pb.site_generator("New", energy=0.4)
+    def site_gen():
+        x = [10, 20, 30]
+        y = [1, 2, 3]
+        z = [0, -1, -2]
+        return x, y, z
+
+    model = pb.Model(graphene.monolayer(), graphene.hexagon_ac(1), site_gen)
+    s = model.system[model.system.x >= 10]
+    assert s.num_sites == 3
+    assert pytest.fuzzy_equal(s.x, [10, 20, 30])
+    assert pytest.fuzzy_equal(s.y, [1, 2, 3])
+    assert pytest.fuzzy_equal(s.z, [0, -1, -2])
+
+
 def test_hopping_generator():
     """Generated next-nearest hoppings should produce the same result as the builtin lattice"""
     from scipy.spatial import cKDTree