implicit.h

Go to the documentation of this file.

/*
    Philote C++ Bindings
 
    Copyright 2022-2025 Christopher A. Lupp
 
    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
    You may obtain a copy of the License at
 
        http://www.apache.org/licenses/LICENSE-2.0
 
    Unless required by applicable law or agreed to in writing, software
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
    limitations under the License.
 
    This work has been cleared for public release, distribution unlimited, case
    number: AFRL-2023-5716.
 
    The views expressed are those of the authors and do not reflect the
    official guidance or position of the United States Government, the
    Department of Defense or of the United States Air Force.
 
    Statement from DoD: The Appearance of external hyperlinks does not
    constitute endorsement by the United States Department of Defense (DoD) of
    the linked websites, of the information, products, or services contained
    therein. The DoD does not exercise any editorial, security, or other
    control over the information you may find at these locations.
*/
#pragma once
 
#include <disciplines.grpc.pb.h>
#include "discipline_server.h"
 
#include <discipline.h>
#include "discipline_client.h"
 
namespace philote
{
    // forward declaration
    class ImplicitDiscipline;
 
    class ImplicitServer : public ImplicitService::Service
    {
    public:
        ImplicitServer() = default;
 
        ~ImplicitServer() noexcept;
 
        void LinkPointers(std::shared_ptr<philote::ImplicitDiscipline> implementation);
 
        void UnlinkPointers();
 
        grpc::Status ComputeResiduals(grpc::ServerContext *context,
                                      grpc::ServerReaderWriter<::philote::Array,
                                                               ::philote::Array> *stream);
 
        grpc::Status SolveResiduals(grpc::ServerContext *context,
                                    grpc::ServerReaderWriter<::philote::Array,
                                                             ::philote::Array> *stream);
 
        grpc::Status ComputeResidualGradients(grpc::ServerContext *context,
                                              grpc::ServerReaderWriter<::philote::Array,
                                                                       ::philote::Array> *stream);
 
        // Test helper methods that accept interface pointers for unit testing with mocks
        template<typename StreamType>
        grpc::Status ComputeResidualsImpl(grpc::ServerContext *context, StreamType *stream);
 
        template<typename StreamType>
        grpc::Status SolveResidualsImpl(grpc::ServerContext *context, StreamType *stream);
 
        template<typename StreamType>
        grpc::Status ComputeResidualGradientsImpl(grpc::ServerContext *context, StreamType *stream);
 
        // Public wrappers for tests
        grpc::Status ComputeResidualsForTesting(grpc::ServerContext *context,
                                                grpc::ServerReaderWriterInterface<::philote::Array,
                                                                                 ::philote::Array> *stream) {
            return ComputeResidualsImpl(context, stream);
        }
 
        grpc::Status SolveResidualsForTesting(grpc::ServerContext *context,
                                              grpc::ServerReaderWriterInterface<::philote::Array,
                                                                               ::philote::Array> *stream) {
            return SolveResidualsImpl(context, stream);
        }
 
        grpc::Status ComputeResidualGradientsForTesting(grpc::ServerContext *context,
                                                        grpc::ServerReaderWriterInterface<::philote::Array,
                                                                                         ::philote::Array> *stream) {
            return ComputeResidualGradientsImpl(context, stream);
        }
 
    private:
        std::shared_ptr<philote::ImplicitDiscipline> implementation_;
    };
 
    class ImplicitDiscipline : public Discipline
    {
    public:
        ImplicitDiscipline();
 
        ~ImplicitDiscipline() noexcept;
 
        void RegisterServices(grpc::ServerBuilder &builder);
 
        void DeclarePartials(const std::string &f, const std::string &x);
 
        virtual void ComputeResiduals(const philote::Variables &inputs,
                                      const philote::Variables &outputs,
                                      philote::Variables &residuals);
 
        virtual void SolveResiduals(const philote::Variables &inputs, philote::Variables &outputs);
 
        virtual void ComputeResidualGradients(const philote::Variables &inputs,
                                              const philote::Variables &outputs,
                                              Partials &partials);
 
    private:
        philote::ImplicitServer implicit_;
        philote::DisciplineServer discipline_server_;
    };
 
    class ImplicitClient : public ::philote::DisciplineClient
    {
    public:
        ImplicitClient() = default;
 
        ~ImplicitClient() noexcept = default;
 
        void ConnectChannel(std::shared_ptr<grpc::ChannelInterface> channel);
 
        Variables ComputeResiduals(const Variables &vars);
 
        Variables SolveResiduals(const Variables &vars);
 
        Partials ComputeResidualGradients(const Variables &vars);
 
        void SetStub(std::unique_ptr<ImplicitService::StubInterface> stub)
        {
            stub_ = std::move(stub);
        }
 
    private:
        std::unique_ptr<ImplicitService::StubInterface> stub_;
    };
 
// Template implementations must be in header
#include <algorithm>
#include <unordered_map>
 
template<typename StreamType>
grpc::Status philote::ImplicitServer::ComputeResidualsImpl(grpc::ServerContext *context, StreamType *stream)
{
    if (!implementation_)
    {
        return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Discipline implementation not linked");
    }
 
    if (!stream)
    {
        return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "Stream is null");
    }
 
    // Check for cancellation before starting
    if (context && context->IsCancelled())
    {
        return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled before start");
    }
 
    philote::Array array;
 
    // preallocate the variables based on meta data
    Variables inputs, outputs, residuals;
    const auto *discipline = static_cast<philote::Discipline *>(implementation_.get());
    if (!discipline)
    {
        return grpc::Status(grpc::StatusCode::INTERNAL, "Failed to cast implementation to Discipline");
    }
 
    for (const auto &var : discipline->var_meta())
    {
        std::string name = var.name();
        if (var.type() == kInput)
            inputs[name] = Variable(var);
        if (var.type() == kOutput)
        {
            outputs[var.name()] = Variable(var);
            residuals[var.name()] = Variable(var);
        }
    }
 
    // Build O(1) lookup map for variable metadata
    std::unordered_map<std::string, const VariableMetaData*> var_lookup;
    for (const auto &var : discipline->var_meta())
    {
        var_lookup[var.name()] = &var;
    }
 
    while (stream->Read(&array))
    {
        // get variables from the stream message
        const std::string &name = array.name();
 
        // get the variable corresponding to the current message using O(1) lookup
        auto var_it = var_lookup.find(name);
        if (var_it == var_lookup.end())
        {
            return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "Variable not found: " + name);
        }
        const VariableMetaData* var = var_it->second;
 
        // Validate that the message type matches the metadata type
        if (array.type() != var->type())
        {
            return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                         "Type mismatch for variable " + name + ": expected " +
                         std::to_string(var->type()) + " but received " + std::to_string(array.type()));
        }
 
        // obtain the inputs and outputs from the stream
        if (var->type() == VariableType::kInput)
        {
            try
            {
                inputs[name].AssignChunk(array);
            }
            catch (const std::exception &e)
            {
                return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                              "Failed to assign chunk for input " + name + ": " + e.what());
            }
        }
        else if (var->type() == VariableType::kOutput)
        {
            try
            {
                outputs[name].AssignChunk(array);
            }
            catch (const std::exception &e)
            {
                return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                              "Failed to assign chunk for output " + name + ": " + e.what());
            }
        }
        else
        {
            return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                         "Invalid variable type received for variable: " + name);
        }
    }
 
    // Check for cancellation before expensive computation
    if (context && context->IsCancelled())
    {
        return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled before computation");
    }
 
    // Set context for discipline to check cancellation during compute
    discipline->SetContext(context);
 
    // call the discipline developer-defined Compute function
    try
    {
        implementation_->ComputeResiduals(inputs, outputs, residuals);
    }
    catch (const std::exception &e)
    {
        discipline->ClearContext();
        return grpc::Status(grpc::StatusCode::INTERNAL,
                      "Failed to compute residuals: " + std::string(e.what()));
    }
 
    // Clear context after computation
    discipline->ClearContext();
 
    // Check for cancellation before sending results
    if (context && context->IsCancelled())
    {
        return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled before sending results");
    }
 
    // iterate through residuals
    for (const auto &res : residuals)
    {
        const std::string &name = res.first;
        try
        {
            res.second.Send(name, "", stream, discipline->stream_opts().num_double(), context);
        }
        catch (const std::exception &e)
        {
            return grpc::Status(grpc::StatusCode::INTERNAL,
                          "Failed to send residual " + name + ": " + e.what());
        }
    }
 
    return grpc::Status::OK;
}
 
template<typename StreamType>
grpc::Status philote::ImplicitServer::SolveResidualsImpl(grpc::ServerContext *context, StreamType *stream)
{
    if (!implementation_)
    {
        return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Discipline implementation not linked");
    }
 
    if (!stream)
    {
        return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "Stream is null");
    }
 
    // Check for cancellation before starting
    if (context && context->IsCancelled())
    {
        return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled before start");
    }
 
    philote::Array array;
 
    // preallocate the inputs based on meta data
    Variables inputs;
    const auto *discipline = static_cast<philote::Discipline *>(implementation_.get());
    if (!discipline)
    {
        return grpc::Status(grpc::StatusCode::INTERNAL, "Failed to cast implementation to Discipline");
    }
 
    for (const auto &var : discipline->var_meta())
    {
        std::string name = var.name();
        if (var.type() == kInput)
            inputs[name] = Variable(var);
    }
 
    // Build O(1) lookup map for variable metadata
    std::unordered_map<std::string, const VariableMetaData*> var_lookup;
    for (const auto &var : discipline->var_meta())
    {
        var_lookup[var.name()] = &var;
    }
 
    while (stream->Read(&array))
    {
        // get variables from the stream message
        const std::string &name = array.name();
 
        // get the variable corresponding to the current message using O(1) lookup
        auto var_it = var_lookup.find(name);
        if (var_it == var_lookup.end())
        {
            return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "Variable not found: " + name);
        }
        const VariableMetaData* var = var_it->second;
 
        // obtain the inputs from the stream (only inputs expected for solve)
        if (var->type() == VariableType::kInput)
        {
            try
            {
                inputs[name].AssignChunk(array);
            }
            catch (const std::exception &e)
            {
                return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                              "Failed to assign chunk for input " + name + ": " + e.what());
            }
        }
        else
        {
            return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                         "Expected input variable but received different type for: " + name);
        }
    }
 
    // preallocate outputs
    Variables outputs;
    for (const VariableMetaData &var : discipline->var_meta())
    {
        if (var.type() == kOutput)
            outputs[var.name()] = Variable(var);
    }
 
    // Check for cancellation before expensive computation
    if (context && context->IsCancelled())
    {
        return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled before computation");
    }
 
    // Set context for discipline to check cancellation during solve
    discipline->SetContext(context);
 
    // call the discipline developer-defined Solve function
    try
    {
        implementation_->SolveResiduals(inputs, outputs);
    }
    catch (const std::exception &e)
    {
        discipline->ClearContext();
        return grpc::Status(grpc::StatusCode::INTERNAL,
                      "Failed to solve residuals: " + std::string(e.what()));
    }
 
    // Clear context after computation
    discipline->ClearContext();
 
    // Check for cancellation before sending results
    if (context && context->IsCancelled())
    {
        return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled before sending results");
    }
 
    // iterate through continuous outputs
    for (const auto &var : outputs)
    {
        const std::string &name = var.first;
        try
        {
            var.second.Send(name, "", stream, discipline->stream_opts().num_double(), context);
        }
        catch (const std::exception &e)
        {
            return grpc::Status(grpc::StatusCode::INTERNAL,
                          "Failed to send output " + name + ": " + e.what());
        }
    }
 
    return grpc::Status::OK;
}
 
template<typename StreamType>
grpc::Status philote::ImplicitServer::ComputeResidualGradientsImpl(grpc::ServerContext *context, StreamType *stream)
{
    if (!implementation_)
    {
        return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Discipline implementation not linked");
    }
 
    if (!stream)
    {
        return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "Stream is null");
    }
 
    // Check for cancellation before starting
    if (context && context->IsCancelled())
    {
        return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled before start");
    }
 
    philote::Array array;
 
    // preallocate the inputs and outputs based on meta data
    Variables inputs, outputs;
    const auto *discipline = static_cast<philote::Discipline *>(implementation_.get());
    if (!discipline)
    {
        return grpc::Status(grpc::StatusCode::INTERNAL, "Failed to cast implementation to Discipline");
    }
 
    for (const auto &var : discipline->var_meta())
    {
        const std::string &name = var.name();
        if (var.type() == kInput)
            inputs[name] = Variable(var);
        if (var.type() == kOutput)
            outputs[name] = Variable(var);
    }
 
    // Build O(1) lookup map for variable metadata
    std::unordered_map<std::string, const VariableMetaData*> var_lookup;
    for (const auto &var : discipline->var_meta())
    {
        var_lookup[var.name()] = &var;
    }
 
    while (stream->Read(&array))
    {
        // get variables from the stream message
        const std::string &name = array.name();
 
        // get the variable corresponding to the current message using O(1) lookup
        auto var_it = var_lookup.find(name);
        if (var_it == var_lookup.end())
        {
            return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "Variable not found: " + name);
        }
        const VariableMetaData* var = var_it->second;
 
        // Validate that the message type matches the metadata type
        if (array.type() != var->type())
        {
            return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                         "Type mismatch for variable " + name + ": expected " +
                         std::to_string(var->type()) + " but received " + std::to_string(array.type()));
        }
 
        // obtain the inputs and outputs from the stream
        if (var->type() == VariableType::kInput)
        {
            try
            {
                inputs[name].AssignChunk(array);
            }
            catch (const std::exception &e)
            {
                return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                              "Failed to assign chunk for input " + name + ": " + e.what());
            }
        }
        else if (var->type() == VariableType::kOutput)
        {
            try
            {
                outputs[name].AssignChunk(array);
            }
            catch (const std::exception &e)
            {
                return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                              "Failed to assign chunk for output " + name + ": " + e.what());
            }
        }
        else
        {
            return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                         "Invalid variable type received for variable: " + name);
        }
    }
 
    // preallocate partials
    Partials partials;
    for (const PartialsMetaData &par : discipline->partials_meta())
    {
        std::vector<size_t> shape;
        for (const int64_t &dim : par.shape())
            shape.push_back(dim);
 
        partials[std::make_pair(par.name(), par.subname())] = Variable(kOutput, shape);
    }
 
    // Check for cancellation before expensive computation
    if (context && context->IsCancelled())
    {
        return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled before computation");
    }
 
    // Set context for discipline to check cancellation during compute
    discipline->SetContext(context);
 
    // call the discipline developer-defined Compute function
    try
    {
        implementation_->ComputeResidualGradients(inputs, outputs, partials);
    }
    catch (const std::exception &e)
    {
        discipline->ClearContext();
        return grpc::Status(grpc::StatusCode::INTERNAL,
                      "Failed to compute residual gradients: " + std::string(e.what()));
    }
 
    // Clear context after computation
    discipline->ClearContext();
 
    // Check for cancellation before sending results
    if (context && context->IsCancelled())
    {
        return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled before sending results");
    }
 
    // iterate through partials
    for (const auto &par : partials)
    {
        const std::string &name = par.first.first;
        const std::string &subname = par.first.second;
        try
        {
            par.second.Send(name, subname, stream, discipline->stream_opts().num_double(), context);
        }
        catch (const std::exception &e)
        {
            return grpc::Status(grpc::StatusCode::INTERNAL,
                          "Failed to send partial " + name + "/" + subname + ": " + e.what());
        }
    }
 
    return grpc::Status::OK;
}} // namespace philote