ENH: Add RPC support

.github/workflows/build_test.yml

@@ -8,7 +8,7 @@ on:
       - main
 
 jobs:
-  buildmamba:
+  buildpixi:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3

.gitignore

@@ -3,6 +3,7 @@
 *.tmp
 subprojects/*
 !subprojects/*.wrap
+*.capnp.*
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

.pre-commit-config.yaml

@@ -26,3 +26,10 @@ repos:
       language: system
       types_or: [c++, c]
       args: ["--filter=-whitespace/comments,-runtime/references,-whitespace/indent,-runtime/int,-build/c++11"]
+    - id: meson-format
+      name: meson format
+      entry: meson
+      language: system
+      types: [meson]
+      files: meson\.build$
+      args: ["format", "-i"]

CppCore/meson.build

@@ -16,6 +16,11 @@ potlib = library(
     install: true,
 )
 
+# ------------------------ Server
+if get_option('with_rpc')
+    subdir('rgpot/rpc')
+endif
+
 # ------------------------ Examples
 
 if get_option('with_examples')

CppCore/rgpot/CuH2/meson.build

@@ -1,7 +1,9 @@
-cuh2 = library('cuh2',
-           'CuH2Pot.cc',
-           'cuh2Utils.cc',
-           dependencies: _deps,
-           cpp_args: _args,
-           include_directories: _incdirs,
-           install : true)
+cuh2 = library(
+    'cuh2',
+    'CuH2Pot.cc',
+    'cuh2Utils.cc',
+    dependencies: _deps,
+    cpp_args: _args,
+    include_directories: _incdirs,
+    install: true,
+)

CppCore/rgpot/LennardJones/meson.build

@@ -1,6 +1,8 @@
-lennard_jones = library('lennard_jones',
-           'LJPot.cc',
-           dependencies: _deps,
-           cpp_args: _args,
-           include_directories: ['../../'],
-           install : true)
+lennard_jones = library(
+    'lennard_jones',
+    'LJPot.cc',
+    dependencies: _deps,
+    cpp_args: _args,
+    include_directories: ['../../'],
+    install: true,
+)

CppCore/rgpot/rpc/Potentials.capnp

@@ -0,0 +1,21 @@
+@0xbd1f89fa17369103;
+
+# Design kanged from eOn v4 C-style structs
+# TODO(rg): Should be more object oriented
+
+struct ForceInput {
+ natm   @0 :Int32; # TODO(rg): Do we really need this..
+ pos    @1 :List(Float64);
+ atmnrs @2 :List(Int32);
+ box    @3 :List(Float64);
+}
+
+struct PotentialResult {
+  energy @0: Float64;
+  forces @1: List(Float64);
+}
+
+interface Potential {
+  calculate @0 (fip :ForceInput)
+    -> (result :PotentialResult);
+}

CppCore/rgpot/rpc/meson.build

@@ -0,0 +1,42 @@
+# Find dependencies and tools
+capnp_rpc_dep = dependency('capnp-rpc')
+capnpc = find_program('capnpc')
+
+output_dir = meson.current_source_dir() / 'gen'
+
+gen_rpc = custom_target(
+    'gen_rpc',
+    input: 'Potentials.capnp',
+    output: ['Potentials.capnp.c++', 'Potentials.capnp.h'],
+    command: [
+        'sh',
+        '-c',
+        'mkdir -p @0@ && @1@ -o c++:@0@ --src-prefix=@2@ @3@'.format(
+            output_dir,
+            capnpc.path(),
+            meson.current_source_dir(),
+            '@INPUT@',
+        ),
+    ],
+)
+
+gen_inc = include_directories('gen')
+_incdirs += gen_inc
+
+ptlrpc = library(
+    'ptlrpc',
+    gen_rpc,
+    dependencies: _deps + capnp_rpc_dep,
+    cpp_args: _args,
+    include_directories: _incdirs,
+    install: true,
+)
+server = executable(
+    'potserv',
+    ['server.cpp', gen_rpc],
+    link_with: [ptlrpc, _linkto],
+    dependencies: _deps + capnp_rpc_dep,
+    cpp_args: _args,
+    include_directories: _incdirs,
+    install: true,
+)

CppCore/rgpot/rpc/pyclient.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python3
+
+import argparse
+import asyncio
+
+import capnp
+import Potentials_capnp
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        usage="Connects to the Calculator server at the given address and does some RPCs"
+    )
+    parser.add_argument("host", help="HOST:PORT")
+
+    return parser.parse_args()
+
+
+async def main(connection):
+    client = capnp.TwoPartyClient(connection)
+    calculator = client.bootstrap().cast_as(Potentials_capnp.Potential)
+
+    print("Connection made ", end="")
+
+    # Create the atomic positions, atom types, and box matrix
+    positions = [
+        0.63940268750835,
+        0.90484742551374,
+        6.97516498544584,  # Cu
+        3.19652040936288,
+        0.90417430354811,
+        6.97547796369474,  # Cu
+        8.98363230369760,
+        9.94703496017833,
+        7.83556854923689,  # H
+        7.64080177576300,
+        9.94703114803832,
+        7.83556986121272,  # H
+    ]
+
+    atom_types = [29, 29, 1, 1]  # Atomic types for Cu and H
+
+    box_matrix = [15.0, 0.0, 0.0, 0.0, 20.0, 0.0, 0.0, 0.0, 30.0]
+
+    # Create the ForceInput message
+    force_input = Potentials_capnp.ForceInput.new_message()
+    force_input.natm = len(atom_types)  # Number of atoms
+
+    # Set positions
+    pos_list = force_input.init("pos", len(positions))
+    for idx, pos in enumerate(positions):
+        pos_list[idx] = pos
+
+    # Set atom types
+    atom_list = force_input.init("atmnrs", len(atom_types))
+    for idx, atom in enumerate(atom_types):
+        atom_list[idx] = atom
+
+    # Set the box matrix
+    box_list = force_input.init("box", len(box_matrix))
+    for idx, val in enumerate(box_matrix):
+        box_list[idx] = val
+
+    # Call the CuH2Pot.calculate method with ForceInput
+    response = calculator.calculate(force_input)
+
+    # Await and print the result
+    result = await response
+    result_dict = result.to_dict()
+    print("Energy:", result_dict["result"]["energy"])
+    print("Forces:", result_dict["result"]["forces"])
+
+
+async def cmd_main(host):
+    host, port = host.split(":")
+    await main(await capnp.AsyncIoStream.create_connection(host=host, port=port))
+
+
+if __name__ == "__main__":
+    asyncio.run(capnp.run(cmd_main(parse_args().host)))

CppCore/rgpot/rpc/server.cpp

@@ -0,0 +1,93 @@
+// MIT License
+// Copyright 2023--present Rohit Goswami <HaoZeke>
+#include <capnp/ez-rpc.h>
+#include <capnp/message.h>
+
+#include "rgpot/CuH2/CuH2Pot.hpp"
+#include "rgpot/LennardJones/LJPot.hpp"
+#include "rgpot/types/AtomMatrix.hpp"
+#include "rgpot/types/adapters/capnp/capnp_adapter.hpp"
+
+class GenericPotImpl final : public Potential::Server {
+private:
+  std::unique_ptr<rgpot::Potential> m_potential;
+
+public:
+  // The constructor takes ownership of a potential object
+  GenericPotImpl(std::unique_ptr<rgpot::Potential> pot)
+      : m_potential(std::move(pot)) {}
+
+  kj::Promise<void> calculate(CalculateContext context) override {
+    // --- Use the adapter to convert FROM Cap'n Proto ---
+    auto fip = context.getParams().getFip();
+    const auto numAtoms = fip.getNatm();
+
+    rgpot::types::AtomMatrix nativePositions =
+        rgpot::types::adapt::capnp::convertPositionsFromCapnp(fip.getPos(),
+                                                              numAtoms);
+    std::vector<int> nativeAtomTypes =
+        rgpot::types::adapt::capnp::convertAtomNumbersFromCapnp(
+            fip.getAtmnrs());
+    std::array<std::array<double, 3>, 3> nativeBoxMatrix =
+        rgpot::types::adapt::capnp::convertBoxMatrixFromCapnp(fip.getBox());
+
+    // --- Call the potential via the polymorphic interface ---
+    // This is the key change: no hardcoded type!
+    auto [energy, forces] =
+        (*m_potential)(nativePositions, nativeAtomTypes, nativeBoxMatrix);
+
+    // --- Use the adapter to populate TO Cap'n Proto ---
+    auto result = context.getResults();
+    auto pres = result.initResult();
+    pres.setEnergy(energy);
+
+    auto forcesList = pres.initForces(numAtoms * 3);
+    rgpot::types::adapt::capnp::populateForcesToCapnp(forcesList, forces);
+
+    return kj::READY_NOW;
+  }
+};
+
+// Main server setup now acts as a factory
+int main(int argc, char *argv[]) {
+  if (argc < 3) {
+    std::cerr << "Usage: " << argv[0] << " <port> <PotentialType>" << std::endl;
+    std::cerr << "  Available PotentialTypes: CuH2"
+              << std::endl; // Add more as you implement them
+    return 1;
+  }
+
+  int port = 12345;
+  try {
+    port = std::stoi(argv[1]);
+  } catch (const std::exception &e) {
+    std::cerr << "Invalid port argument '" << argv[1]
+              << "'. Using default 12345." << std::endl;
+  }
+
+  std::string pot_type = argv[2];
+  std::unique_ptr<rgpot::Potential> potential_to_use;
+
+  if (pot_type == "CuH2") {
+    std::cout << "Loading CuH2 potential..." << std::endl;
+    potential_to_use = std::make_unique<rgpot::CuH2Pot>();
+  } else if (pot_type == "LJ") {
+    std::cout << "Loading LJ potential..." << std::endl;
+    potential_to_use = std::make_unique<rgpot::LJPot>();
+  } else {
+    std::cerr << "Error: Unknown potential type '" << pot_type << "'"
+              << std::endl;
+    return 1;
+  }
+
+  capnp::EzRpcServer server(
+      kj::heap<GenericPotImpl>(std::move(potential_to_use)), "localhost", port);
+
+  // Keep the server running indefinitely
+  auto &waitScope = server.getWaitScope();
+  std::cout << "Server running on port " << port << " with " << pot_type
+            << " potential." << std::endl;
+  kj::NEVER_DONE.wait(waitScope);
+
+  return 0;
+}