write_params.py

#!/usr/bin/python
# Copyright 2020 Makani Technologies LLC
#
# 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.


r"""Writes parameters out to various file formats.

This function is the main point of entry for generating parameter
files.  Using the Makani configuration system (mconfig), this function
creates the parameters for a specified wing model and writes the
output to a specified file format (e.g. JSON or a .c file).

E.g.:
  ./config/write_params.py \
    --wing_model=m600 \
    --type=json \
    --input_file=config/m600/system_params.py \
    --output_file=m600.json

For other examples, see control/BUILD and sim/BUILD.
"""

import json
import os
import re
import sys
import textwrap

import gflags
from makani.config import mconfig
from makani.config import overrides_util
from makani.lib.python import wing_flag

wing_flag.AppeaseLintWhenImportingFlagOnly()
FLAGS = gflags.FLAGS

gflags.DEFINE_enum('type', 'json',
                   ['json', 'system', 'control', 'sim', 'monitor'],
                   'Specifies the type of the output configuration file.')

gflags.DEFINE_string('input_file', None,
                     'Full path to the input configuration file.',
                     short_name='i')
gflags.DEFINE_string('output_file', None,
                     'Full path to the ouput file.',
                     short_name='o')
gflags.DEFINE_string('overrides', None,
                     'JSON string of overrides.')


def WriteJsonParams(params, filename):
  """Outputs a parameters dictionary as a JSON file.

  Args:
    params: Parameters dictionary to convert to JSON.
    filename: Filename to write JSON to.
  """
  with open(filename, 'w') as f:
    json.dump(params, f, indent=2, separators=(',', ': '))


def WriteDefaultSystemParams(system_params, filename):
  """Writes the system parameters.

  The system parameters are defined by the SystemParams struct in
  control/system_types.h.  The default values for these parameters are
  compiled into the program by initializing a static struct in
  control/system_params.c, which is generated by this function.
  Additionally, this auto-generated file provides pointers to the
  "global" system parameters, which are accessible in any program that
  links to this file.

  Args:
    system_params: System parameters dictionary to convert to a .c file.
    filename: Filename of the .c file.
  """
  with open(filename, 'w') as f:
    f.write(textwrap.dedent("""
        #include "control/system_params.h"
        #include "control/system_types.h"

        static SystemParams system_params = %s;

        const GlobalSystemParams g_sys = {
          .ts = &system_params.ts,
          .phys = &system_params.phys,
          .tether = &system_params.tether,
          .wing = &system_params.wing,
          .perch = &system_params.perch,
          .rotors = {
            %s
          }
        };

        const SystemParams *GetSystemParams(void) { return &system_params; }
        SystemParams *GetSystemParamsUnsafe(void) { return &system_params; }
        """)[1:] % (
            ConvertDictToCStructStr(system_params),
            ',\n    '.join(['&system_params.rotors[%d]' % i
                            for i in range(len(system_params['rotors']))])))


def WriteDefaultControlParams(control_params, filename):
  """Writes the control parameters.

  See the comment for WriteDefaultSystemParams above for more details.

  Args:
    control_params: Control parameters dictionary to convert to a .c file.
    filename: Filename of the .c file.
  """
  with open(filename, 'w') as f:
    f.write(textwrap.dedent("""
        #include "control/control_params.h"
        #include "control/control_types.h"

        static ControlParams control_params = %s;

        GlobalControlParams g_cont = {
          .flight_plan = &control_params.flight_plan,
          .control_opt = &control_params.control_opt,
          .joystick_control = &control_params.joystick_control,
          .simple_aero_model = &control_params.simple_aero_model,
          .rotor_control = &control_params.rotor_control
        };

        const ControlParams *GetControlParams(void) { return &control_params; }
        ControlParams *GetControlParamsUnsafe(void) { return &control_params; }
        """[1:]) % ConvertDictToCStructStr(control_params))


def WriteDefaultSimParams(sim_params, filename):
  """Writes the simulation parameters.

  See the comment for WriteDefaultSystemParams above for more details.

  Args:
    sim_params: Simulator parameters dictionary to convert to a .c file.
    filename: Filename of the .c file.
  """
  with open(filename, 'w') as f:
    f.write(textwrap.dedent("""
        #include "sim/sim_params.h"
        #include "sim/sim_types.h"

        static SimParams sim_params = %s;

        const GlobalSimParams g_sim = {
          .sim_opt = &sim_params.sim_opt,
          .phys_sim = &sim_params.phys_sim
        };

        const SimParams *GetSimParams(void) { return &sim_params; }
        SimParams *GetSimParamsUnsafe(void) { return &sim_params; }
        """[1:]) % ConvertDictToCStructStr(sim_params))


def WriteDefaultMonitorParams(monitor_params, filename):
  """Writes the monitor parameters.

  See the comment for WriteDefaultSystemParams above for more details.

  Args:
    monitor_params: Monitor parameters dictionary to convert to a .c file.
    filename: Filename of the .c file.
  """
  with open(filename, 'w') as f:
    f.write(textwrap.dedent("""
        #include "gs/monitor/monitor_params.h"
        #include "gs/monitor/monitor_types.h"

        static MonitorParams monitor_params = %s;

        const MonitorParams *GetMonitorParams(void) { return &monitor_params; }
        MonitorParams *GetMonitorParamsUnsafe(void) { return &monitor_params; }
        """[1:]) % ConvertDictToCStructStr(monitor_params))


def ConvertDictToCStructStr(d):
  """Converts a dict to a string represent a C struct initialization."""
  # Convert JSON string to C struct string.
  json_str = json.dumps(d, indent=2, separators=(',', ' = '))
  c_struct_str = re.sub(r'"(\w+)" = ', r'.\1 = ', json_str)
  # The spaces in front of the brackets are necessary to avoid
  # converting the brackets in a string describing an array, e.g.
  # motor_connections[0] --> motor_connections{0}.
  c_struct_str = c_struct_str.replace(' [', ' {').replace(' ]', ' }')
  # Makes cpplint.py complain less.
  c_struct_str = re.sub(re.compile(r'\n\s*{', re.MULTILINE), ' {', c_struct_str)
  return c_struct_str


def main(argv):
  # Parse flags.
  try:
    argv = FLAGS(argv)
  except gflags.FlagsError, e:
    print '\nError: %s\n\nUsage: %s ARGS\n%s' % (e, sys.argv[0], FLAGS)
    sys.exit(1)

  # If no input file is specified, revert to the default for the given
  # file type.
  if FLAGS.input_file is None:
    if FLAGS.type == 'json':
      assert False, '--input_file must be specified for --type=json.'
    if FLAGS.type == 'control':
      assert False, '--input_file must be specified for --type=control.'
    elif FLAGS.type == 'monitor':
      assert False, '--input_file must be specified for --type=monitor.'
    elif FLAGS.type == 'sim':
      assert False, '--input_file must be specified for --type=sim.'
    elif FLAGS.type == 'system':
      assert False, '--input_file must be specified for --type=system.'
    else:
      assert False, '--type was not properly validated.'

  # Convert a configuration filename to the equivalent module name by
  # remove the beginning "config/" and the trailing ".py", and then
  # replacing "/" with ".".
  module_str = FLAGS.input_file[7:-3].replace('/', '.')

  # Parse overrides command line option.
  if FLAGS.overrides:
    overrides = overrides_util.PreprocessOverrides(json.loads(FLAGS.overrides))
  else:
    overrides = {}

  # Make a parameters dict from the configuration files.
  params = mconfig.MakeParams(module_str, overrides=overrides,
                              override_method='derived')

  # Write parameters to the output_file or to stdout.
  if FLAGS.output_file is None:
    print params
  else:
    if FLAGS.type == 'json':
      WriteJsonParams(params, FLAGS.output_file)
    elif FLAGS.type == 'control':
      WriteDefaultControlParams(params, FLAGS.output_file)
    elif FLAGS.type == 'monitor':
      WriteDefaultMonitorParams(params, FLAGS.output_file)
    elif FLAGS.type == 'sim':
      WriteDefaultSimParams(params, FLAGS.output_file)
    elif FLAGS.type == 'system':
      WriteDefaultSystemParams(params, FLAGS.output_file)
    else:
      assert False, '--type was not properly validated.'


if __name__ == '__main__':
  gflags.RegisterValidator(
      'input_file',
      lambda f: f is None or os.path.isfile(f),
      message='input_file does not exist.')
  gflags.RegisterValidator(
      'input_file',
      lambda f: f is None or (f.startswith('config/') and f.endswith('.py')),
      message='input_file must be a .py file under the config/ directory.')

  main(sys.argv)