swerve_trajectory_generator.hpp

// Copyright (c) TrajoptLib contributors
 
#pragma once
 
#include <stdint.h>
 
#include <cmath>
#include <cstddef>
#include <expected>
#include <utility>
#include <vector>
 
#include <sleipnir/autodiff/variable.hpp>
#include <sleipnir/optimization/problem.hpp>
#include <sleipnir/optimization/solver/exit_status.hpp>
 
#include "trajopt/geometry/translation2.hpp"
#include "trajopt/path/path.hpp"
#include "trajopt/path/path_builder.hpp"
#include "trajopt/util/symbol_exports.hpp"
 
namespace trajopt {

struct TRAJOPT_DLLEXPORT SwerveDrivetrain {
  double mass;

  double moi;

  double wheel_radius;

  double wheel_max_angular_velocity;

  double wheel_max_torque;

  double wheel_cof;

  std::vector<Translation2d> modules;
};

struct TRAJOPT_DLLEXPORT SwerveSolution {
  std::vector<double> dt;

  std::vector<double> x;

  std::vector<double> y;

  std::vector<double> thetacos;

  std::vector<double> thetasin;

  std::vector<double> vx;

  std::vector<double> vy;

  std::vector<double> omega;

  std::vector<double> ax;

  std::vector<double> ay;

  std::vector<double> alpha;

  std::vector<std::vector<double>> module_fx;

  std::vector<std::vector<double>> module_fy;
};

class TRAJOPT_DLLEXPORT SwerveTrajectorySample {
 public:
  double timestamp = 0.0;

  double x = 0.0;

  double y = 0.0;

  double heading = 0.0;

  double velocity_x = 0.0;

  double velocity_y = 0.0;

  double angular_velocity = 0.0;

  double acceleration_x = 0.0;

  double acceleration_y = 0.0;

  double angular_acceleration = 0.0;

  std::vector<double> module_forces_x;

  std::vector<double> module_forces_y;
 
  SwerveTrajectorySample() = default;

  SwerveTrajectorySample(double timestamp, double x, double y, double heading,
                         double velocity_x, double velocity_y,
                         double angular_velocity, double acceleration_x,
                         double acceleration_y, double angular_acceleration,
                         std::vector<double> module_forces_x,
                         std::vector<double> module_forces_y)
      : timestamp{timestamp},
        x{x},
        y{y},
        heading{heading},
        velocity_x{velocity_x},
        velocity_y{velocity_y},
        angular_velocity{angular_velocity},
        acceleration_x{acceleration_x},
        acceleration_y{acceleration_y},
        angular_acceleration{angular_acceleration},
        module_forces_x{std::move(module_forces_x)},
        module_forces_y{std::move(module_forces_y)} {}
};

class TRAJOPT_DLLEXPORT SwerveTrajectory {
 public:
  std::vector<SwerveTrajectorySample> samples;
 
  SwerveTrajectory() = default;

  explicit SwerveTrajectory(std::vector<SwerveTrajectorySample> samples)
      : samples{std::move(samples)} {}

  explicit SwerveTrajectory(const SwerveSolution& solution) {
    double ts = 0.0;
    for (size_t sample = 0; sample < solution.x.size(); ++sample) {
      samples.emplace_back(
          ts, solution.x[sample], solution.y[sample],
          std::atan2(solution.thetasin[sample], solution.thetacos[sample]),
          solution.vx[sample], solution.vy[sample], solution.omega[sample],
          solution.ax[sample], solution.ay[sample], solution.alpha[sample],
          solution.module_fx[sample], solution.module_fy[sample]);
      ts += solution.dt[sample];
    }
  }
};

using SwervePath = Path<SwerveDrivetrain, SwerveSolution>;

using SwervePathBuilder = PathBuilder<SwerveDrivetrain, SwerveSolution>;

class TRAJOPT_DLLEXPORT SwerveTrajectoryGenerator {
 public:
  explicit SwerveTrajectoryGenerator(SwervePathBuilder path_builder,
                                     int64_t handle = 0);

  std::expected<SwerveSolution, slp::ExitStatus> generate(
      bool diagnostics = false);
 
 private:
  SwervePath path;

  std::vector<slp::Variable<double>> x;
  std::vector<slp::Variable<double>> y;
  std::vector<slp::Variable<double>> cosθ;
  std::vector<slp::Variable<double>> sinθ;
  std::vector<slp::Variable<double>> vx;
  std::vector<slp::Variable<double>> vy;
  std::vector<slp::Variable<double>> ω;
  std::vector<slp::Variable<double>> ax;
  std::vector<slp::Variable<double>> ay;
  std::vector<slp::Variable<double>> α;

  std::vector<std::vector<slp::Variable<double>>> Fx;
  std::vector<std::vector<slp::Variable<double>>> Fy;

  std::vector<slp::Variable<double>> dts;

  std::vector<size_t> Ns;
 
  slp::Problem<double> problem;
 
  void apply_initial_guess(const SwerveSolution& solution);
 
  SwerveSolution construct_swerve_solution();
};
 
}  // namespace trajopt