<commit_before>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

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.
==============================================================================*/
#include <memory>

#include "absl/cleanup/cleanup.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/xla/service/compiler.h"
#include "tensorflow/compiler/xla/service/executable.h"
#include "tensorflow/compiler/xla/service/hlo_cost_analysis.h"
#include "tensorflow/compiler/xla/service/hlo_module.h"
#include "tensorflow/compiler/xla/service/hlo_module_group.h"
#include "tensorflow/compiler/xla/service/shaped_buffer.h"
#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/compiler/xla/stream_executor/device_memory_allocator.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/c_api_conversions.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/c_api_decl.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/proto_helper.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/status_helper.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_executable.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_executor.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_executor_c_api.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_platform.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_platform_id.h"
#include "tensorflow/compiler/xla/util.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"

namespace xla {

namespace {

using ::tensorflow::tpu::ExecutorApiFn;

class TpuCompiler : public Compiler {
 public:
  TpuCompiler() { compiler_ = ExecutorApiFn()->TpuCompiler_NewFn(); }
  ~TpuCompiler() override { ExecutorApiFn()->TpuCompiler_FreeFn(compiler_); }

  stream_executor::Platform::Id PlatformId() const override {
    return tensorflow::tpu::GetTpuPlatformId();
  }

  StatusOr<std::unique_ptr<HloModule>> RunHloPasses(
      std::unique_ptr<HloModule> module,
      stream_executor::StreamExecutor* executor,
      const CompileOptions& options) override {
    XLA_HloModule hlo_module;
    auto cleanup = absl::MakeCleanup([&hlo_module]() {
      stream_executor::tpu::SerializedProto_Free(hlo_module.proto);
      ApiConverter::Destroy(&hlo_module.module_config);
    });
    hlo_module.module_config = ApiConverter::ToC(module->config());
    hlo_module.proto = stream_executor::tpu::SerializeProto(module->ToProto());
    auto allocator = ApiConverter::ToC(options.device_allocator);
    XLA_HloModule result;
    StatusHelper status;
    ExecutorApiFn()->TpuCompiler_RunHloPassesFn(
        compiler_, &hlo_module,
        static_cast<tensorflow::tpu::TpuExecutor*>(executor->implementation())
            ->se_executor(),
        &allocator, &result, status.c_status);
    if (!status.ok()) {
      return status.status();
    }
    HloModuleProto result_proto =
        stream_executor::tpu::DeserializeProto<HloModuleProto>(result.proto);
    stream_executor::tpu::SerializedProto_Free(result.proto);
    return HloModule::CreateFromProto(result_proto, module->config());
  }

  StatusOr<std::unique_ptr<Executable>> RunBackend(
      std::unique_ptr<HloModule> module,
      stream_executor::StreamExecutor* executor,
      const CompileOptions& options) override {
    XLA_HloModule hlo_module;
    auto cleanup = absl::MakeCleanup([&hlo_module]() {
      stream_executor::tpu::SerializedProto_Free(hlo_module.proto);
      ApiConverter::Destroy(&hlo_module.module_config);
    });
    SE_Executable* result;
    hlo_module.module_config = ApiConverter::ToC(module->config());
    hlo_module.proto = stream_executor::tpu::SerializeProto(module->ToProto());
    auto allocator = ApiConverter::ToC(options.device_allocator);

    StatusHelper status;
    ExecutorApiFn()->TpuCompiler_RunBackendFn(
        compiler_, &hlo_module,
        static_cast<tensorflow::tpu::TpuExecutor*>(executor->implementation())
            ->se_executor(),
        &allocator, &result, status.c_status);
    if (!status.ok()) {
      return status.status();
    }

    std::unique_ptr<Executable> exec =
        std::make_unique<TpuExecutable>(result, std::move(module));
    return exec;
  }

  StatusOr<std::vector<std::unique_ptr<Executable>>> Compile(
      std::unique_ptr<HloModuleGroup> module_group,
      std::vector<std::vector<stream_executor::StreamExecutor*>> stream_exec,
      const CompileOptions& options) override {
    XLA_HloModuleGroup se_module_group;
    se_module_group.proto =
        stream_executor::tpu::SerializeProto(module_group->ToProto());
    se_module_group.module_config =
        new XLA_HloModuleConfig[module_group->size()];
    int module_group_size = module_group->size();
    auto cleanup_config =
        absl::MakeCleanup([&se_module_group, module_group_size]() {
          for (auto i = 0; i < module_group_size; ++i) {
            ApiConverter::Destroy(&se_module_group.module_config[i]);
          }
          delete[] se_module_group.module_config;
        });
    for (int i = 0; i < module_group->size(); ++i) {
      const auto& config = module_group->module(i).config();
      se_module_group.module_config[i] = ApiConverter::ToC(config);
    }
    std::vector<SE_StreamExecutorList> se_lists(stream_exec.size());
    std::vector<std::vector<SE_StreamExecutor*>> se_lists_storage;
    for (int i = 0; i < stream_exec.size(); ++i) {
      se_lists[i].count = stream_exec[i].size();
      se_lists_storage.emplace_back(stream_exec[i].size());
      se_lists[i].exec = se_lists_storage.back().data();
      for (int j = 0; j < stream_exec[i].size(); ++j) {
        se_lists[i].exec[j] = static_cast<tensorflow::tpu::TpuExecutor*>(
                                  stream_exec[i][j]->implementation())
                                  ->se_executor();
      }
    }

    SE_DeviceMemoryAllocator allocator =
        ApiConverter::ToC(options.device_allocator);

    SE_Executable** se_executables = new SE_Executable*[module_group->size()];

    StatusHelper status;

    ExecutorApiFn()->TpuCompiler_CompileFn(
        compiler_, &se_module_group, se_lists.data(), stream_exec.size(),
        &allocator, se_executables, status.c_status);

    if (!status.ok()) {
      return status.status();
    }

    std::vector<std::unique_ptr<Executable>> executables;
    for (int i = 0; i < module_group->size(); ++i) {
      // We get the HloModule from the compiled executable, rather than reusing
      // the input module from 'module_group', in case the module changed in
      // some way. For example, if the computation is automatically partitioned
      // via XLA, the executable's module may have different input/output shapes
      // than the input module.
      XLA_HloModule c_module =
          ExecutorApiFn()->TpuExecutable_HloModuleFn(se_executables[i]);
      auto cleanup_c_module = absl::MakeCleanup(
          [&c_module]() { ApiConverter::Destroy(&c_module); });
      TF_ASSIGN_OR_RETURN(std::unique_ptr<HloModule> module,
                          ApiConverter::FromC(c_module));
      std::shared_ptr<HloModule> module_shared(module.release());
      executables.emplace_back(std::make_unique<TpuExecutable>(
          se_executables[i], std::move(module_shared)));
    }

    stream_executor::tpu::SerializedProto_Free(se_module_group.proto);
    delete[] se_executables;

    return executables;
  }

  // Compiles the HLO module group for ahead-of-time execution.  This is
  // intended for use in static compilation.
  StatusOr<std::vector<std::unique_ptr<AotCompilationResult>>>
  CompileAheadOfTime(std::unique_ptr<HloModuleGroup> module_group,
                     const AotCompilationOptions& options) override {
    return Unimplemented("This compiler does not support CompileAheadOfTime.");
  }

  // Returns a function that computes the size in bytes of the logical
  // buffer that contains a shape.
  HloCostAnalysis::ShapeSizeFunction ShapeSizeBytesFunction() const override {
    return [this](const xla::Shape& shape) {
      XLA_Shape c_shape;
      ApiConverter::ToC(shape, &c_shape);
      int64_t bytes =
          ExecutorApiFn()->TpuCompiler_ShapeSizeFn(compiler_, &c_shape);
      ApiConverter::Destroy(&c_shape);
      return bytes;
    };
  }

  Shape DefaultDeviceShapeRepresentation(const Shape& shape) const override {
    XLA_Shape host_shape, device_shape;
    ApiConverter::ToC(shape, &host_shape);
    ExecutorApiFn()->TpuCompiler_DefaultDeviceShapeRepresentationFn(
        compiler_, &host_shape, &device_shape);
    return ApiConverter::FromC(&device_shape);
  }

 private:
  Tpu_Compiler* compiler_;
};

static bool InitModule() {
  xla::Compiler::RegisterCompilerFactory(
      tensorflow::tpu::GetTpuPlatformId(),
      []() { return std::make_unique<TpuCompiler>(); });
  return true;
}

static bool module_initialized = InitModule();

}  // namespace
}  // namespace xla
<commit_msg>[libtpu] Fix some memory leaks<commit_after>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

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.
==============================================================================*/
#include <memory>

#include "absl/cleanup/cleanup.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/xla/service/compiler.h"
#include "tensorflow/compiler/xla/service/executable.h"
#include "tensorflow/compiler/xla/service/hlo_cost_analysis.h"
#include "tensorflow/compiler/xla/service/hlo_module.h"
#include "tensorflow/compiler/xla/service/hlo_module_group.h"
#include "tensorflow/compiler/xla/service/shaped_buffer.h"
#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/compiler/xla/stream_executor/device_memory_allocator.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/c_api_conversions.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/c_api_decl.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/proto_helper.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/status_helper.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_executable.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_executor.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_executor_c_api.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_platform.h"
#include "tensorflow/compiler/xla/stream_executor/tpu/tpu_platform_id.h"
#include "tensorflow/compiler/xla/util.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"

namespace xla {

namespace {

using ::tensorflow::tpu::ExecutorApiFn;

class TpuCompiler : public Compiler {
 public:
  TpuCompiler() { compiler_ = ExecutorApiFn()->TpuCompiler_NewFn(); }
  ~TpuCompiler() override { ExecutorApiFn()->TpuCompiler_FreeFn(compiler_); }

  stream_executor::Platform::Id PlatformId() const override {
    return tensorflow::tpu::GetTpuPlatformId();
  }

  StatusOr<std::unique_ptr<HloModule>> RunHloPasses(
      std::unique_ptr<HloModule> module,
      stream_executor::StreamExecutor* executor,
      const CompileOptions& options) override {
    XLA_HloModule hlo_module;
    auto cleanup = absl::MakeCleanup([&hlo_module]() {
      stream_executor::tpu::SerializedProto_Free(hlo_module.proto);
      ApiConverter::Destroy(&hlo_module.module_config);
    });
    hlo_module.module_config = ApiConverter::ToC(module->config());
    hlo_module.proto = stream_executor::tpu::SerializeProto(module->ToProto());
    auto allocator = ApiConverter::ToC(options.device_allocator);
    XLA_HloModule result;
    StatusHelper status;
    ExecutorApiFn()->TpuCompiler_RunHloPassesFn(
        compiler_, &hlo_module,
        static_cast<tensorflow::tpu::TpuExecutor*>(executor->implementation())
            ->se_executor(),
        &allocator, &result, status.c_status);
    if (!status.ok()) {
      return status.status();
    }
    HloModuleProto result_proto =
        stream_executor::tpu::DeserializeProto<HloModuleProto>(result.proto);
    stream_executor::tpu::SerializedProto_Free(result.proto);
    return HloModule::CreateFromProto(result_proto, module->config());
  }

  StatusOr<std::unique_ptr<Executable>> RunBackend(
      std::unique_ptr<HloModule> module,
      stream_executor::StreamExecutor* executor,
      const CompileOptions& options) override {
    XLA_HloModule hlo_module;
    auto cleanup = absl::MakeCleanup([&hlo_module]() {
      stream_executor::tpu::SerializedProto_Free(hlo_module.proto);
      ApiConverter::Destroy(&hlo_module.module_config);
    });
    SE_Executable* result;
    hlo_module.module_config = ApiConverter::ToC(module->config());
    hlo_module.proto = stream_executor::tpu::SerializeProto(module->ToProto());
    auto allocator = ApiConverter::ToC(options.device_allocator);

    StatusHelper status;
    ExecutorApiFn()->TpuCompiler_RunBackendFn(
        compiler_, &hlo_module,
        static_cast<tensorflow::tpu::TpuExecutor*>(executor->implementation())
            ->se_executor(),
        &allocator, &result, status.c_status);
    if (!status.ok()) {
      return status.status();
    }

    std::unique_ptr<Executable> exec =
        std::make_unique<TpuExecutable>(result, std::move(module));
    return exec;
  }

  StatusOr<std::vector<std::unique_ptr<Executable>>> Compile(
      std::unique_ptr<HloModuleGroup> module_group,
      std::vector<std::vector<stream_executor::StreamExecutor*>> stream_exec,
      const CompileOptions& options) override {
    XLA_HloModuleGroup se_module_group;
    se_module_group.proto =
        stream_executor::tpu::SerializeProto(module_group->ToProto());
    se_module_group.module_config =
        new XLA_HloModuleConfig[module_group->size()];
    int module_group_size = module_group->size();
    auto cleanup_config =
        absl::MakeCleanup([&se_module_group, module_group_size]() {
          for (auto i = 0; i < module_group_size; ++i) {
            ApiConverter::Destroy(&se_module_group.module_config[i]);
          }
          delete[] se_module_group.module_config;
        });
    for (int i = 0; i < module_group->size(); ++i) {
      const auto& config = module_group->module(i).config();
      se_module_group.module_config[i] = ApiConverter::ToC(config);
    }
    std::vector<SE_StreamExecutorList> se_lists(stream_exec.size());
    std::vector<std::vector<SE_StreamExecutor*>> se_lists_storage;
    for (int i = 0; i < stream_exec.size(); ++i) {
      se_lists[i].count = stream_exec[i].size();
      se_lists_storage.emplace_back(stream_exec[i].size());
      se_lists[i].exec = se_lists_storage.back().data();
      for (int j = 0; j < stream_exec[i].size(); ++j) {
        se_lists[i].exec[j] = static_cast<tensorflow::tpu::TpuExecutor*>(
                                  stream_exec[i][j]->implementation())
                                  ->se_executor();
      }
    }

    SE_DeviceMemoryAllocator allocator =
        ApiConverter::ToC(options.device_allocator);

    SE_Executable** se_executables = new SE_Executable*[module_group->size()];

    StatusHelper status;

    ExecutorApiFn()->TpuCompiler_CompileFn(
        compiler_, &se_module_group, se_lists.data(), stream_exec.size(),
        &allocator, se_executables, status.c_status);

    if (!status.ok()) {
      return status.status();
    }

    std::vector<std::unique_ptr<Executable>> executables;
    for (int i = 0; i < module_group->size(); ++i) {
      // We get the HloModule from the compiled executable, rather than reusing
      // the input module from 'module_group', in case the module changed in
      // some way. For example, if the computation is automatically partitioned
      // via XLA, the executable's module may have different input/output shapes
      // than the input module.
      XLA_HloModule c_module =
          ExecutorApiFn()->TpuExecutable_HloModuleFn(se_executables[i]);
      auto cleanup_c_module = absl::MakeCleanup(
          [&c_module]() { ApiConverter::Destroy(&c_module); });
      TF_ASSIGN_OR_RETURN(std::unique_ptr<HloModule> module,
                          ApiConverter::FromC(c_module));
      std::shared_ptr<HloModule> module_shared(module.release());
      executables.emplace_back(std::make_unique<TpuExecutable>(
          se_executables[i], std::move(module_shared)));
    }

    stream_executor::tpu::SerializedProto_Free(se_module_group.proto);
    delete[] se_executables;

    return executables;
  }

  // Compiles the HLO module group for ahead-of-time execution.  This is
  // intended for use in static compilation.
  StatusOr<std::vector<std::unique_ptr<AotCompilationResult>>>
  CompileAheadOfTime(std::unique_ptr<HloModuleGroup> module_group,
                     const AotCompilationOptions& options) override {
    return Unimplemented("This compiler does not support CompileAheadOfTime.");
  }

  // Returns a function that computes the size in bytes of the logical
  // buffer that contains a shape.
  HloCostAnalysis::ShapeSizeFunction ShapeSizeBytesFunction() const override {
    return [this](const xla::Shape& shape) {
      XLA_Shape c_shape;
      ApiConverter::ToC(shape, &c_shape);
      int64_t bytes =
          ExecutorApiFn()->TpuCompiler_ShapeSizeFn(compiler_, &c_shape);
      ApiConverter::Destroy(&c_shape);
      return bytes;
    };
  }

  Shape DefaultDeviceShapeRepresentation(const Shape& shape) const override {
    XLA_Shape host_shape, device_shape;
    ApiConverter::ToC(shape, &host_shape);
    ExecutorApiFn()->TpuCompiler_DefaultDeviceShapeRepresentationFn(
        compiler_, &host_shape, &device_shape);
    ApiConverter::Destroy(&host_shape);
    Shape result = ApiConverter::FromC(&device_shape);
    ApiConverter::Destroy(&device_shape);
    return result;
  }

 private:
  Tpu_Compiler* compiler_;
};

static bool InitModule() {
  xla::Compiler::RegisterCompilerFactory(
      tensorflow::tpu::GetTpuPlatformId(),
      []() { return std::make_unique<TpuCompiler>(); });
  return true;
}

static bool module_initialized = InitModule();

}  // namespace
}  // namespace xla
<|endoftext|>


<commit_before>#include "MainForm.moc"

#include <QtGui>
#include <iostream>
#include <sstream>
#include <tbb/tbb.h>
#include <png.h>

MainForm::MainForm(QGraphicsScene* img, nbt* bf, QWidget* parent_)
                 : QGraphicsView(img, parent_), scene_(), bf_(bf), scale_(1),
                   images() {
  connect(this, SIGNAL(scaleSig()), this, SLOT(scale()));
  connect(this, SIGNAL(renderNewImage()), this, SLOT(populateSceneItem()));
  connect(this, SIGNAL(saveToFileSignal()), this, SLOT(saveToFile()));
  setTransformationAnchor(QGraphicsView::AnchorUnderMouse);
}

class ApplyFoo {
  MainForm* mainform_;
  int i_;
  tbb::atomic<int>* index_;
 public:
  void operator()( const tbb::blocked_range<int32_t>& r ) const {
    for(int32_t j=r.begin(); j!=r.end(); ++j) {
      bool result = false;
      QPoint bp = mainform_->projectCoords(QPoint(j, i_),
                                        (4 - mainform_->bf_->set().rotate) % 4);
      const Image<uint8_t>& image = mainform_->bf_-> getImage(bp.x(),
                                                              bp.y(), &result);
      if (!result) {
        continue;
      }
      *index_ += 1;
      mainform_->images.push(MainForm::image_coords(image, QPoint(bp.x(),
                                                                  bp.y())));
      mainform_->renderNewImageEmitter();
    }
  }
  ApplyFoo(MainForm* mainform, int i, tbb::atomic<int>* index)
          : mainform_(mainform), i_(i), index_(index) {}
  /* just for the compiler */
  ApplyFoo(const ApplyFoo& rhs)
          : mainform_(rhs.mainform_), i_(rhs.i_), index_(rhs.index_) {}
 private:
  ApplyFoo& operator=(const ApplyFoo&);
};

void MainForm::populateScene() {
  tbb::atomic<int> index;
  index = 0;
  QPoint min(bf_->xPos_min(), bf_->zPos_min());
  min = projectCoords(min, bf_->set().rotate);
  QPoint max(bf_->xPos_max(), bf_->zPos_max());
  max = projectCoords(max, bf_->set().rotate);
  min_norm = QPoint(std::min(min.x(), max.x()),
                    std::min(min.y(), max.y()));
  max_norm = QPoint(std::max(min.x(), max.x()),
                    std::max(min.y(), max.y()));

    std::stringstream ss;
    int width = (max_norm.x() - min_norm.x() + 1) * 16;
    int height = (max_norm.y() - min_norm.y() + 1) * 16;
    if (bf_->set().oblique) height += 128;
    ss << "P7\n"
       << "WIDTH "     << width << "\n"
       << "HEIGHT "    << height << "\n"
       << "DEPTH "     << 4 << "\n"
       << "MAXVAL "    << 255 << "\n"
       << "TUPLTYPE "  << "RGB_ALPHA" << "\n"
       << "ENDHDR"     << "\n";
    header_size = ss.str().size();
    FILE* pam = fopen("test.ppm", "w");
    fwrite(ss.str().c_str(), 1, header_size, pam);
    fseek(pam, width * height * 4 - 1, SEEK_CUR);
    fwrite("", 1, 1, pam);
    fclose(pam);

  for (int i = min_norm.y(); i <= max_norm.y(); ++i) {
    tbb::parallel_for(tbb::blocked_range<int32_t>(min_norm.x(),
                                                  max_norm.x() + 1),
                                                  ApplyFoo(this, i, &index));
    if (index > 10000) {
      std::cerr << "cache cleared!" << std::endl;
      index = 0;
      bf_->clearCache();
    }
  }
  bf_->clearCache();
  emit saveToFileSignal();

  // QPen pen;
  // pen.setColor(QColor(255, 0, 0, 255));
  // scene()->addEllipse(185, 50, 5, 5, pen);
  // rotate(270.0f);
  // 187 52
  // setTransform(QTransform().scale(1, 1));

  size_t nr_pixels = width * height;
  pam = fopen("test.ppm", "r");
  FILE* out = fopen("test.png", "w");
  fseek(pam, header_size, SEEK_CUR);
  png_struct* pngP;
  png_info* infoP;
  pngP = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
  infoP = png_create_info_struct(pngP);
  png_set_IHDR(pngP, infoP, width, height, 8,
               PNG_COLOR_TYPE_RGB_ALPHA,
               PNG_INTERLACE_NONE,
               PNG_COMPRESSION_TYPE_DEFAULT,
               PNG_FILTER_TYPE_DEFAULT);
  png_init_io(pngP, out);
  png_write_info(pngP, infoP);

  png_byte* pngRow = reinterpret_cast<png_byte*>(malloc(width * 4));
  for (int i = 0; i < height; ++i) {
    fread(pngRow, 4, width, pam);
    png_write_row(pngP, pngRow);
  }
  free(pngRow);

  png_write_end(pngP, infoP);
  png_destroy_write_struct(&pngP, &infoP);
  fclose(pam);
  fclose(out);
}

void MainForm::renderNewImageEmitter() {
  emit renderNewImage();
}

void MainForm::saveToFile() {
  scale_ = 1;
  scale();
  QImage image(mapFromScene(scene()->sceneRect()).boundingRect().adjusted(0, 0, -1, -1).size(), QImage::Format_ARGB32);
  image.fill(0);
  QPainter painter(&image);
  render(&painter, painter.viewport(), mapFromScene(scene()->sceneRect()).boundingRect().adjusted(0, 0, -1, -1));
  // image.save("image.png");
  // exit(1);
}

QPoint MainForm::projectCoords(QPoint p, int phi) {
  return projectCoords(p.x(), p.y(), phi);
}

QPoint MainForm::projectCoords(int _x, int _y, int phi) {
  if (phi == 0) return QPoint(_x, _y);
  int cos_phi = phi % 2 - 1;
  if (!phi) cos_phi = 1;
  int sin_phi = phi % 2;
  if (phi == 3) sin_phi = -1;
  QPoint ret;
  ret.setX(_x * cos_phi - _y * sin_phi);
  ret.setY(_x * sin_phi + _y * cos_phi);
  return ret;
}

void MainForm::populateSceneItem() {
  MainForm::image_coords img_coor;
  if (images.try_pop(img_coor)) {
    QImage img(&(img_coor.first.data[0]),
               img_coor.first.cols,
               img_coor.first.rows,
               QImage::Format_ARGB32);
    QGraphicsPixmapItem* pi = scene()->addPixmap(QPixmap::fromImage(img));
    pi->setFlag(QGraphicsItem::ItemIsMovable, false);
    pi->setFlag(QGraphicsItem::ItemIsSelectable, false);
    QPoint projected = projectCoords(16 * img_coor.second.x(),
                                     16 * img_coor.second.y(),
                                     bf_->set().rotate);
    pi->setPos(projected);
    if (bf_->set().rotate == 0) {
      pi->setZValue(img_coor.second.y());
    } else if (bf_->set().rotate == 1) {
      pi->setZValue(img_coor.second.x());
    } else if (bf_->set().rotate == 2) {
      pi->setZValue(-img_coor.second.y());
    } else if (bf_->set().rotate == 3) {
      pi->setZValue(-img_coor.second.x());
    }

    size_t width = img_coor.first.cols;
    size_t height = img_coor.first.rows;
    size_t nr_pixels = width * height;
    int offset_x = projected.x() - min_norm.x() * 16;
    int offset_y = projected.y() - min_norm.y() * 16;
    int g_width = (max_norm.x() - min_norm.x() + 1) * 16;
    int g_height = (max_norm.y() - min_norm.y() + 1) * 16;

    FILE* pam = fopen("test.ppm", "r+");
    fseek(pam, header_size, SEEK_CUR);
    for (size_t i = 0; i < nr_pixels; ++i) {
      size_t index = i * 4;
      std::swap(img_coor.first.data[index], img_coor.first.data[index + 2]);
    }
    fseek(pam, offset_y * g_width * 4 + offset_x * 4, SEEK_CUR);
    for (size_t i = 0; i < height; ++i) {
      for (size_t j = 0; j < width; ++j) {
        if (img_coor.first.data[i * width * 4 + j * 4 + 3] != 0) {
          fwrite(&(img_coor.first.data[i * width * 4 + j * 4]), 4, 1, pam);
        } else {
          fseek(pam, 4, SEEK_CUR);
        }
      }
      fseek(pam, g_width * 4 - width * 4, SEEK_CUR);
    }
    fclose(pam);
  } else {
    std::cerr << "must not happen!" << std::endl;
    exit(1);
  }
}

void MainForm::scale() {
  if (scale_ <= 0) {
    setTransform(QTransform().scale(1.0 / pow(2.0, abs(scale_ - 1)),
                                    1.0 / pow(2.0, abs(scale_ - 1))));
  } else {
    setTransform(QTransform().scale(scale_, scale_));
  }
}

void MainForm::mousePressEvent(QMouseEvent* mevent) {
  switch (mevent->button()) {
    case Qt::LeftButton:
      ++scale_;
      emit scaleSig();
      break;
    case Qt::MidButton:
      break;
    case Qt::RightButton:
      --scale_;
      emit scaleSig();
      break;
    default:
      break;
  }
  return;
}

void MainForm::mouseDoubleClickEvent(QMouseEvent* mevent) {
  mousePressEvent(mevent);
  return;
}
<commit_msg>render scene not view, remove png/pam stuff from MainForm<commit_after>#include "MainForm.moc"

#include <QtGui>
#include <iostream>
#include <sstream>
#include <tbb/tbb.h>
#include <png.h>

MainForm::MainForm(QGraphicsScene* img, nbt* bf, QWidget* parent_)
                 : QGraphicsView(img, parent_), scene_(), bf_(bf), scale_(1),
                   images() {
  connect(this, SIGNAL(scaleSig()), this, SLOT(scale()));
  connect(this, SIGNAL(renderNewImage()), this, SLOT(populateSceneItem()));
  connect(this, SIGNAL(saveToFileSignal()), this, SLOT(saveToFile()));
  setTransformationAnchor(QGraphicsView::AnchorUnderMouse);
}

class ApplyFoo {
  MainForm* mainform_;
  int i_;
  tbb::atomic<int>* index_;
 public:
  void operator()( const tbb::blocked_range<int32_t>& r ) const {
    for(int32_t j=r.begin(); j!=r.end(); ++j) {
      bool result = false;
      QPoint bp = mainform_->projectCoords(QPoint(j, i_),
                                        (4 - mainform_->bf_->set().rotate) % 4);
      const Image<uint8_t>& image = mainform_->bf_-> getImage(bp.x(),
                                                              bp.y(), &result);
      if (!result) {
        continue;
      }
      *index_ += 1;
      mainform_->images.push(MainForm::image_coords(image, QPoint(bp.x(),
                                                                  bp.y())));
      mainform_->renderNewImageEmitter();
    }
  }
  ApplyFoo(MainForm* mainform, int i, tbb::atomic<int>* index)
          : mainform_(mainform), i_(i), index_(index) {}
  /* just for the compiler */
  ApplyFoo(const ApplyFoo& rhs)
          : mainform_(rhs.mainform_), i_(rhs.i_), index_(rhs.index_) {}
 private:
  ApplyFoo& operator=(const ApplyFoo&);
};

void MainForm::populateScene() {
  tbb::atomic<int> index;
  index = 0;
  QPoint min(bf_->xPos_min(), bf_->zPos_min());
  min = projectCoords(min, bf_->set().rotate);
  QPoint max(bf_->xPos_max(), bf_->zPos_max());
  max = projectCoords(max, bf_->set().rotate);
  min_norm = QPoint(std::min(min.x(), max.x()),
                    std::min(min.y(), max.y()));
  max_norm = QPoint(std::max(min.x(), max.x()),
                    std::max(min.y(), max.y()));

  for (int i = min_norm.y(); i <= max_norm.y(); ++i) {
    tbb::parallel_for(tbb::blocked_range<int32_t>(min_norm.x(),
                                                  max_norm.x() + 1),
                                                  ApplyFoo(this, i, &index));
    if (index > 10000) {
      std::cerr << "cache cleared!" << std::endl;
      index = 0;
      bf_->clearCache();
    }
  }
  bf_->clearCache();
  emit saveToFileSignal();
}

void MainForm::renderNewImageEmitter() {
  emit renderNewImage();
}

void MainForm::saveToFile() {
  QImage image(scene()->sceneRect().toRect().size(), QImage::Format_ARGB32);
  image.fill(0);
  QPainter painter(&image);
  scene()->render(&painter, painter.viewport(), scene()->sceneRect());
  image.save("image.png");
  fprintf(stderr, "image saved!\n");
  // exit(1);
}

QPoint MainForm::projectCoords(QPoint p, int phi) {
  return projectCoords(p.x(), p.y(), phi);
}

QPoint MainForm::projectCoords(int _x, int _y, int phi) {
  if (phi == 0) return QPoint(_x, _y);
  int cos_phi = phi % 2 - 1;
  if (!phi) cos_phi = 1;
  int sin_phi = phi % 2;
  if (phi == 3) sin_phi = -1;
  QPoint ret;
  ret.setX(_x * cos_phi - _y * sin_phi);
  ret.setY(_x * sin_phi + _y * cos_phi);
  return ret;
}

void MainForm::populateSceneItem() {
  MainForm::image_coords img_coor;
  if (images.try_pop(img_coor)) {
    QImage img(&(img_coor.first.data[0]),
               img_coor.first.cols,
               img_coor.first.rows,
               QImage::Format_ARGB32);
    QGraphicsPixmapItem* pi = scene()->addPixmap(QPixmap::fromImage(img));
    pi->setFlag(QGraphicsItem::ItemIsMovable, false);
    pi->setFlag(QGraphicsItem::ItemIsSelectable, false);
    QPoint projected = projectCoords(16 * img_coor.second.x(),
                                     16 * img_coor.second.y(),
                                     bf_->set().rotate);
    pi->setPos(projected);
    if (bf_->set().rotate == 0) {
      pi->setZValue(img_coor.second.y());
    } else if (bf_->set().rotate == 1) {
      pi->setZValue(img_coor.second.x());
    } else if (bf_->set().rotate == 2) {
      pi->setZValue(-img_coor.second.y());
    } else if (bf_->set().rotate == 3) {
      pi->setZValue(-img_coor.second.x());
    }
  } else {
    std::cerr << "must not happen!" << std::endl;
    exit(1);
  }
}

void MainForm::scale() {
  if (scale_ <= 0) {
    setTransform(QTransform().scale(1.0 / pow(2.0, abs(scale_ - 1)),
                                    1.0 / pow(2.0, abs(scale_ - 1))));
  } else {
    setTransform(QTransform().scale(scale_, scale_));
  }
}

void MainForm::mousePressEvent(QMouseEvent* mevent) {
  switch (mevent->button()) {
    case Qt::LeftButton:
      ++scale_;
      emit scaleSig();
      break;
    case Qt::MidButton:
      break;
    case Qt::RightButton:
      --scale_;
      emit scaleSig();
      break;
    default:
      break;
  }
  return;
}

void MainForm::mouseDoubleClickEvent(QMouseEvent* mevent) {
  mousePressEvent(mevent);
  return;
}
<|endoftext|>


<commit_before>//* This file is part of the MOOSE framework
//* https://www.mooseframework.org
//*
//* All rights reserved, see COPYRIGHT for full restrictions
//* https://github.com/idaholab/moose/blob/master/COPYRIGHT
//*
//* Licensed under LGPL 2.1, please see LICENSE for details
//* https://www.gnu.org/licenses/lgpl-2.1.html

#include "MortarGapHeatTransferAction.h"

#include "AddVariableAction.h"
#include "FEProblem.h"
#include "libmesh/string_to_enum.h"
#include "NonlinearSystem.h"

#include "ModularGapConductanceConstraint.h"
#include "GapFluxModelRadiation.h"
#include "GapFluxModelConduction.h"

// Counter for modular user objects
static unsigned int thermal_action_userobject_radiation_counter = 0;
static unsigned int thermal_action_userobject_conduction_counter = 0;

registerMooseAction("HeatConductionApp", MortarGapHeatTransferAction, "append_mesh_generator");
registerMooseAction("HeatConductionApp", MortarGapHeatTransferAction, "add_mortar_variable");
registerMooseAction("HeatConductionApp", MortarGapHeatTransferAction, "add_constraint");
registerMooseAction("HeatConductionApp", MortarGapHeatTransferAction, "add_user_object");

InputParameters
MortarGapHeatTransferAction::validParams()
{
  InputParameters params = Action::validParams();
  params.addClassDescription(
      "Action that controls the creation of all of the necessary objects for "
      "calculation of heat transfer through an open/closed gap");

  // Modular mortar gap conductance
  MooseEnum gap_heat_transfer_formulation("point_segment mortar", "mortar");
  params.addParam<MooseEnum>(
      "formulation", gap_heat_transfer_formulation, "The gap heat transfer formulation to be used");

  params.addParam<Real>("thermal_lm_scaling",
                        1.,
                        "Scaling factor to apply to the thermal Lagrange multiplier variable");

  params += ModularGapConductanceConstraint::validParams();
  params += GapFluxModelRadiation::validParams();
  params += GapFluxModelConduction::validParams();
  //
  params.makeParamNotRequired<std::vector<BoundaryName>>("boundary");
  params.makeParamNotRequired<BoundaryName>("primary_boundary");
  params.makeParamNotRequired<SubdomainName>("primary_subdomain");
  params.makeParamNotRequired<BoundaryName>("secondary_boundary");
  params.makeParamNotRequired<SubdomainName>("secondary_subdomain");
  params.makeParamNotRequired<std::vector<VariableName>>("temperature");

  params.addParam<std::vector<UserObjectName>>(
      "user_object_physics",
      "The list of physical contributions to gap heat transfer implemented in user objects");

  return params;
}

MortarGapHeatTransferAction::MortarGapHeatTransferAction(const InputParameters & params)
  : Action(params),
    _formulation(getParam<MooseEnum>("formulation").getEnum<GapHeatTransferFormulation>()),
    _user_provided_mortar_meshes(false)
{

  if (params.isParamSetByUser("primary_subdomain") &&
      params.isParamSetByUser("secondary_subdomain"))
  {
    mooseInfo("Mortar gap heat transfer action is using the lower-dimensional domains provided by "
              "the user");
    _user_provided_mortar_meshes = true;
  }
  else
    mooseInfo("Mortar gap heat transfer action is creating new lower-dimensional domains");
}

void
MortarGapHeatTransferAction::act()
{
  if (_formulation == GapHeatTransferFormulation::MORTAR)
  {

    if (_current_task == "append_mesh_generator")
      addMortarMesh();
    else if (_current_task == "add_mortar_variable")
      addMortarVariable();
    if (_current_task == "add_constraint")
      addConstraints();
    else if (_current_task == "add_user_object")
      addUserObjects();
  }
  else
    paramError(
        "formulation",
        "The formulation selected to solve gap heat transfer physics is not currently available.");
}

void
MortarGapHeatTransferAction::coreMortarMesh()
{
  if (!(_app.isRecovering() && _app.isUltimateMaster()) && !_app.masterMesh())
  {
    std::string action_name = MooseUtils::shortName(name());

    const MeshGeneratorName primary_name = action_name + "_primary_subdomain" + "_generator";
    const MeshGeneratorName secondary_name = action_name + "_secondary_subdomain" + "_generator";

    auto primary_params = _factory.getValidParams("LowerDBlockFromSidesetGenerator");
    auto secondary_params = _factory.getValidParams("LowerDBlockFromSidesetGenerator");

    primary_params.set<SubdomainName>("new_block_name") = action_name + "_primary_subdomain";
    secondary_params.set<SubdomainName>("new_block_name") = action_name + "_secondary_subdomain";

    primary_params.set<std::vector<BoundaryName>>("sidesets") = {
        getParam<BoundaryName>("primary_boundary")};
    secondary_params.set<std::vector<BoundaryName>>("sidesets") = {
        getParam<BoundaryName>("secondary_boundary")};

    _app.appendMeshGenerator("LowerDBlockFromSidesetGenerator", primary_name, primary_params);
    _app.appendMeshGenerator("LowerDBlockFromSidesetGenerator", secondary_name, secondary_params);
  }
}

void
MortarGapHeatTransferAction::addRelationshipManagers(Moose::RelationshipManagerType input_rm_type)
{
  checkForExistingSubdomains();
  std::string action_name = MooseUtils::shortName(name());

  auto params = MortarConstraintBase::validParams();
  params.set<bool>("use_displaced_mesh") = getParam<bool>("use_displaced_mesh");
  params.set<BoundaryName>("primary_boundary") = getParam<BoundaryName>("primary_boundary");
  params.set<BoundaryName>("secondary_boundary") = getParam<BoundaryName>("secondary_boundary");

  if (_user_provided_mortar_meshes)
  {
    params.set<SubdomainName>("primary_subdomain") = getParam<SubdomainName>("primary_subdomain");
    params.set<SubdomainName>("secondary_subdomain") =
        getParam<SubdomainName>("secondary_subdomain");
  }
  else
  {
    params.set<SubdomainName>("primary_subdomain") = action_name + "_primary_subdomain";
    params.set<SubdomainName>("secondary_subdomain") = action_name + "_secondary_subdomain";
  }

  addRelationshipManagers(input_rm_type, params);
}

void
MortarGapHeatTransferAction::addMortarVariable()
{
  checkForExistingSubdomains();

  InputParameters params = _factory.getValidParams("MooseVariableBase");

  const std::string & temperature = getParam<std::vector<VariableName>>("temperature")[0];
  std::string action_name = MooseUtils::shortName(name());

  mooseAssert(_problem->systemBaseNonlinear().hasVariable(temperature),
              "Temperature variable is missing");

  const auto primal_type = _problem->systemBaseNonlinear().system().variable_type(temperature);
  const int lm_order = primal_type.order.get_order();

  if (primal_type.family != LAGRANGE)
    mooseError("The mortar thermal action can only be used with LAGRANGE finite elements");

  params.set<MooseEnum>("family") = Utility::enum_to_string<FEFamily>(primal_type.family);
  params.set<MooseEnum>("order") = Utility::enum_to_string<Order>(OrderWrapper{lm_order});

  if (_user_provided_mortar_meshes)
    params.set<std::vector<SubdomainName>>("block") = {
        getParam<SubdomainName>("secondary_subdomain")};
  else
    params.set<std::vector<SubdomainName>>("block") = {action_name + "_secondary_subdomain"};

  params.set<std::vector<Real>>("scaling") = {getParam<Real>("thermal_lm_scaling")};
  auto fe_type = AddVariableAction::feType(params);
  auto var_type = AddVariableAction::variableType(fe_type);

  _problem->addVariable(var_type, action_name + "_thermal_lm", params);
}

void
MortarGapHeatTransferAction::addConstraints()
{
  checkForExistingSubdomains();

  InputParameters params = _factory.getValidParams("ModularGapConductanceConstraint");
  const std::string action_name = MooseUtils::shortName(name());

  params.applyParameters(parameters());
  params.set<bool>("correct_edge_dropping") = true;
  params.set<bool>("use_displaced_mesh") = true;

  params.set<BoundaryName>("primary_boundary") = getParam<BoundaryName>("primary_boundary");
  params.set<BoundaryName>("secondary_boundary") = getParam<BoundaryName>("secondary_boundary");

  if (_user_provided_mortar_meshes)
  {
    params.set<SubdomainName>("primary_subdomain") = getParam<SubdomainName>("primary_subdomain");
    params.set<SubdomainName>("secondary_subdomain") =
        getParam<SubdomainName>("secondary_subdomain");
  }
  else
  {
    params.set<SubdomainName>("primary_subdomain") = action_name + "_primary_subdomain";
    params.set<SubdomainName>("secondary_subdomain") = action_name + "_secondary_subdomain";
  }

  params.set<NonlinearVariableName>("variable") = action_name + "_thermal_lm";

  params.set<VariableName>("secondary_variable") =
      getParam<std::vector<VariableName>>("temperature")[0];

  const std::vector<UserObjectName> user_object_vector =
      getParam<std::vector<UserObjectName>>("user_object_physics");

  std::vector<UserObjectName> uoname_strings(0);
  unsigned int conduction_index = 0;
  unsigned int radiation_index = 0;

  for (const auto i : index_range(user_object_vector))
  {
    if (user_object_vector[i] == "GapFluxModelConduction")
      uoname_strings.push_back("gap_flux_model_conduction_object_" + MooseUtils::shortName(name()) +
                               "_" + Moose::stringify(conduction_index++));
    else if (user_object_vector[i] == "GapFluxModelRadiation")
      uoname_strings.push_back("gap_flux_model_radiation_object_" + MooseUtils::shortName(name()) +
                               "_" + Moose::stringify(radiation_index++));
  }

  params.set<std::vector<UserObjectName>>("gap_flux_models") = uoname_strings;
  _problem->addConstraint(
      "ModularGapConductanceConstraint", action_name + "_ModularGapConductanceConstraint", params);
}

void
MortarGapHeatTransferAction::addMortarMesh()
{
  // Let's browse over existing mechanical actions to see if the primary and secondary
  // subdomains have been created
  checkForExistingSubdomains();

  // We may have available lower-dimensional domains (e.g. from a mechanical contact action), whose
  // subdomains can be reused for adding mortar variables and constraints.
  if (!_user_provided_mortar_meshes)
    coreMortarMesh();
}

void
MortarGapHeatTransferAction::addUserObjects()
{
  // It is risky to apply this optimization to contact problems
  // since the problem configuration may be changed during Jacobian
  // evaluation. We therefore turn it off for all contact problems so that
  // PETSc-3.8.4 or higher will have the same behavior as PETSc-3.8.3 or older.
  mooseAssert(_problem, "Problem pointer is null");

  if (!_problem->isSNESMFReuseBaseSetbyUser())
    _problem->setSNESMFReuseBase(false, false);

  const std::vector<UserObjectName> user_object_vector =
      getParam<std::vector<UserObjectName>>("user_object_physics");

  for (const auto i : index_range(user_object_vector))
  {
    if (user_object_vector[i] == "GapFluxModelConduction")
    {
      auto var_params = _factory.getValidParams("GapFluxModelConduction");

      var_params.set<std::vector<VariableName>>("temperature") =
          getParam<std::vector<VariableName>>("temperature");
      var_params.set<Real>("gap_conductivity") = getParam<Real>("gap_conductivity");

      if (isParamValid("gap_conductivity_function"))
        var_params.set<FunctionName>("gap_conductivity_function") =
            getParam<FunctionName>("gap_conductivity_function");

      if (isParamValid("gap_conductivity_function_variable"))
        var_params.set<std::vector<VariableName>>("gap_conductivity_function_variable") =
            getParam<std::vector<VariableName>>("gap_conductivity_function_variable");

      var_params.set<Real>("min_gap") = getParam<Real>("min_gap");
      var_params.set<unsigned int>("min_gap_order") = getParam<unsigned int>("min_gap_order");

      var_params.set<std::vector<BoundaryName>>("boundary") =
          getParam<std::vector<BoundaryName>>("boundary");

      var_params.set<bool>("use_displaced_mesh") = true;

      _problem->addUserObject("GapFluxModelConduction",
                              "gap_flux_model_conduction_object_" + MooseUtils::shortName(name()) +
                                  "_" +
                                  Moose::stringify(thermal_action_userobject_conduction_counter++),
                              var_params);
    }
    else if (user_object_vector[i] == "GapFluxModelRadiation")
    {
      auto var_params = _factory.getValidParams("GapFluxModelRadiation");

      var_params.set<Real>("stefan_boltzmann") = getParam<Real>("stefan_boltzmann");
      var_params.set<Real>("primary_emissivity") = getParam<Real>("primary_emissivity");
      var_params.set<Real>("secondary_emissivity") = getParam<Real>("secondary_emissivity");

      var_params.set<std::vector<BoundaryName>>("boundary") =
          getParam<std::vector<BoundaryName>>("boundary");

      var_params.set<std::vector<VariableName>>("temperature") =
          getParam<std::vector<VariableName>>("temperature");

      var_params.set<bool>("use_displaced_mesh") = true;

      _problem->addUserObject("GapFluxModelRadiation",
                              "gap_flux_model_radiation_object_" + MooseUtils::shortName(name()) +
                                  "_" +
                                  Moose::stringify(thermal_action_userobject_radiation_counter++),
                              var_params);
    }
    else
      paramError("user_object_physics",
                 "At least one of the user objects provided to capture gap heat transfer physics "
                 "using a mortar formulation are not currently supported");
  }
}

void
MortarGapHeatTransferAction::checkForExistingSubdomains()
{
  if (parameters().isParamSetByUser("primary_subdomain") &&
      parameters().isParamSetByUser("secondary_subdomain"))
  {
    _user_provided_mortar_meshes = true;
  }
  else
    _user_provided_mortar_meshes = false;
}
<commit_msg>Only subdomains are not required (can be created in the action)<commit_after>//* This file is part of the MOOSE framework
//* https://www.mooseframework.org
//*
//* All rights reserved, see COPYRIGHT for full restrictions
//* https://github.com/idaholab/moose/blob/master/COPYRIGHT
//*
//* Licensed under LGPL 2.1, please see LICENSE for details
//* https://www.gnu.org/licenses/lgpl-2.1.html

#include "MortarGapHeatTransferAction.h"

#include "AddVariableAction.h"
#include "FEProblem.h"
#include "libmesh/string_to_enum.h"
#include "NonlinearSystem.h"

#include "ModularGapConductanceConstraint.h"
#include "GapFluxModelRadiation.h"
#include "GapFluxModelConduction.h"

// Counter for modular user objects
static unsigned int thermal_action_userobject_radiation_counter = 0;
static unsigned int thermal_action_userobject_conduction_counter = 0;

registerMooseAction("HeatConductionApp", MortarGapHeatTransferAction, "append_mesh_generator");
registerMooseAction("HeatConductionApp", MortarGapHeatTransferAction, "add_mortar_variable");
registerMooseAction("HeatConductionApp", MortarGapHeatTransferAction, "add_constraint");
registerMooseAction("HeatConductionApp", MortarGapHeatTransferAction, "add_user_object");

InputParameters
MortarGapHeatTransferAction::validParams()
{
  InputParameters params = Action::validParams();
  params.addClassDescription(
      "Action that controls the creation of all of the necessary objects for "
      "calculation of heat transfer through an open/closed gap");

  // Modular mortar gap conductance
  MooseEnum gap_heat_transfer_formulation("point_segment mortar", "mortar");
  params.addParam<MooseEnum>(
      "formulation", gap_heat_transfer_formulation, "The gap heat transfer formulation to be used");

  params.addParam<Real>("thermal_lm_scaling",
                        1.,
                        "Scaling factor to apply to the thermal Lagrange multiplier variable");

  params += ModularGapConductanceConstraint::validParams();
  params += GapFluxModelRadiation::validParams();
  params += GapFluxModelConduction::validParams();

  params.makeParamNotRequired<SubdomainName>("primary_subdomain");
  params.makeParamNotRequired<SubdomainName>("secondary_subdomain");

  params.addParam<std::vector<UserObjectName>>(
      "user_object_physics",
      "The list of physical contributions to gap heat transfer implemented in user objects");

  return params;
}

MortarGapHeatTransferAction::MortarGapHeatTransferAction(const InputParameters & params)
  : Action(params),
    _formulation(getParam<MooseEnum>("formulation").getEnum<GapHeatTransferFormulation>()),
    _user_provided_mortar_meshes(false)
{

  if (params.isParamSetByUser("primary_subdomain") &&
      params.isParamSetByUser("secondary_subdomain"))
  {
    mooseInfo("Mortar gap heat transfer action is using the lower-dimensional domains provided by "
              "the user");
    _user_provided_mortar_meshes = true;
  }
  else
    mooseInfo("Mortar gap heat transfer action is creating new lower-dimensional domains");
}

void
MortarGapHeatTransferAction::act()
{
  if (_formulation == GapHeatTransferFormulation::MORTAR)
  {

    if (_current_task == "append_mesh_generator")
      addMortarMesh();
    else if (_current_task == "add_mortar_variable")
      addMortarVariable();
    if (_current_task == "add_constraint")
      addConstraints();
    else if (_current_task == "add_user_object")
      addUserObjects();
  }
  else
    paramError(
        "formulation",
        "The formulation selected to solve gap heat transfer physics is not currently available.");
}

void
MortarGapHeatTransferAction::coreMortarMesh()
{
  if (!(_app.isRecovering() && _app.isUltimateMaster()) && !_app.masterMesh())
  {
    std::string action_name = MooseUtils::shortName(name());

    const MeshGeneratorName primary_name = action_name + "_primary_subdomain" + "_generator";
    const MeshGeneratorName secondary_name = action_name + "_secondary_subdomain" + "_generator";

    auto primary_params = _factory.getValidParams("LowerDBlockFromSidesetGenerator");
    auto secondary_params = _factory.getValidParams("LowerDBlockFromSidesetGenerator");

    primary_params.set<SubdomainName>("new_block_name") = action_name + "_primary_subdomain";
    secondary_params.set<SubdomainName>("new_block_name") = action_name + "_secondary_subdomain";

    primary_params.set<std::vector<BoundaryName>>("sidesets") = {
        getParam<BoundaryName>("primary_boundary")};
    secondary_params.set<std::vector<BoundaryName>>("sidesets") = {
        getParam<BoundaryName>("secondary_boundary")};

    _app.appendMeshGenerator("LowerDBlockFromSidesetGenerator", primary_name, primary_params);
    _app.appendMeshGenerator("LowerDBlockFromSidesetGenerator", secondary_name, secondary_params);
  }
}

void
MortarGapHeatTransferAction::addRelationshipManagers(Moose::RelationshipManagerType input_rm_type)
{
  checkForExistingSubdomains();
  std::string action_name = MooseUtils::shortName(name());

  auto params = MortarConstraintBase::validParams();
  params.set<bool>("use_displaced_mesh") = getParam<bool>("use_displaced_mesh");
  params.set<BoundaryName>("primary_boundary") = getParam<BoundaryName>("primary_boundary");
  params.set<BoundaryName>("secondary_boundary") = getParam<BoundaryName>("secondary_boundary");

  if (_user_provided_mortar_meshes)
  {
    params.set<SubdomainName>("primary_subdomain") = getParam<SubdomainName>("primary_subdomain");
    params.set<SubdomainName>("secondary_subdomain") =
        getParam<SubdomainName>("secondary_subdomain");
  }
  else
  {
    params.set<SubdomainName>("primary_subdomain") = action_name + "_primary_subdomain";
    params.set<SubdomainName>("secondary_subdomain") = action_name + "_secondary_subdomain";
  }

  addRelationshipManagers(input_rm_type, params);
}

void
MortarGapHeatTransferAction::addMortarVariable()
{
  checkForExistingSubdomains();

  InputParameters params = _factory.getValidParams("MooseVariableBase");

  const std::string & temperature = getParam<std::vector<VariableName>>("temperature")[0];
  std::string action_name = MooseUtils::shortName(name());

  mooseAssert(_problem->systemBaseNonlinear().hasVariable(temperature),
              "Temperature variable is missing");

  const auto primal_type = _problem->systemBaseNonlinear().system().variable_type(temperature);
  const int lm_order = primal_type.order.get_order();

  if (primal_type.family != LAGRANGE)
    mooseError("The mortar thermal action can only be used with LAGRANGE finite elements");

  params.set<MooseEnum>("family") = Utility::enum_to_string<FEFamily>(primal_type.family);
  params.set<MooseEnum>("order") = Utility::enum_to_string<Order>(OrderWrapper{lm_order});

  if (_user_provided_mortar_meshes)
    params.set<std::vector<SubdomainName>>("block") = {
        getParam<SubdomainName>("secondary_subdomain")};
  else
    params.set<std::vector<SubdomainName>>("block") = {action_name + "_secondary_subdomain"};

  params.set<std::vector<Real>>("scaling") = {getParam<Real>("thermal_lm_scaling")};
  auto fe_type = AddVariableAction::feType(params);
  auto var_type = AddVariableAction::variableType(fe_type);

  _problem->addVariable(var_type, action_name + "_thermal_lm", params);
}

void
MortarGapHeatTransferAction::addConstraints()
{
  checkForExistingSubdomains();

  InputParameters params = _factory.getValidParams("ModularGapConductanceConstraint");
  const std::string action_name = MooseUtils::shortName(name());

  params.applyParameters(parameters());
  params.set<bool>("correct_edge_dropping") = true;
  params.set<bool>("use_displaced_mesh") = true;

  params.set<BoundaryName>("primary_boundary") = getParam<BoundaryName>("primary_boundary");
  params.set<BoundaryName>("secondary_boundary") = getParam<BoundaryName>("secondary_boundary");

  if (_user_provided_mortar_meshes)
  {
    params.set<SubdomainName>("primary_subdomain") = getParam<SubdomainName>("primary_subdomain");
    params.set<SubdomainName>("secondary_subdomain") =
        getParam<SubdomainName>("secondary_subdomain");
  }
  else
  {
    params.set<SubdomainName>("primary_subdomain") = action_name + "_primary_subdomain";
    params.set<SubdomainName>("secondary_subdomain") = action_name + "_secondary_subdomain";
  }

  params.set<NonlinearVariableName>("variable") = action_name + "_thermal_lm";

  params.set<VariableName>("secondary_variable") =
      getParam<std::vector<VariableName>>("temperature")[0];

  const std::vector<UserObjectName> user_object_vector =
      getParam<std::vector<UserObjectName>>("user_object_physics");

  std::vector<UserObjectName> uoname_strings(0);
  unsigned int conduction_index = 0;
  unsigned int radiation_index = 0;

  for (const auto i : index_range(user_object_vector))
  {
    if (user_object_vector[i] == "GapFluxModelConduction")
      uoname_strings.push_back("gap_flux_model_conduction_object_" + MooseUtils::shortName(name()) +
                               "_" + Moose::stringify(conduction_index++));
    else if (user_object_vector[i] == "GapFluxModelRadiation")
      uoname_strings.push_back("gap_flux_model_radiation_object_" + MooseUtils::shortName(name()) +
                               "_" + Moose::stringify(radiation_index++));
  }

  params.set<std::vector<UserObjectName>>("gap_flux_models") = uoname_strings;
  _problem->addConstraint(
      "ModularGapConductanceConstraint", action_name + "_ModularGapConductanceConstraint", params);
}

void
MortarGapHeatTransferAction::addMortarMesh()
{
  // Let's browse over existing mechanical actions to see if the primary and secondary
  // subdomains have been created
  checkForExistingSubdomains();

  // We may have available lower-dimensional domains (e.g. from a mechanical contact action), whose
  // subdomains can be reused for adding mortar variables and constraints.
  if (!_user_provided_mortar_meshes)
    coreMortarMesh();
}

void
MortarGapHeatTransferAction::addUserObjects()
{
  // It is risky to apply this optimization to contact problems
  // since the problem configuration may be changed during Jacobian
  // evaluation. We therefore turn it off for all contact problems so that
  // PETSc-3.8.4 or higher will have the same behavior as PETSc-3.8.3 or older.
  mooseAssert(_problem, "Problem pointer is null");

  if (!_problem->isSNESMFReuseBaseSetbyUser())
    _problem->setSNESMFReuseBase(false, false);

  const std::vector<UserObjectName> user_object_vector =
      getParam<std::vector<UserObjectName>>("user_object_physics");

  for (const auto i : index_range(user_object_vector))
  {
    if (user_object_vector[i] == "GapFluxModelConduction")
    {
      auto var_params = _factory.getValidParams("GapFluxModelConduction");

      var_params.set<std::vector<VariableName>>("temperature") =
          getParam<std::vector<VariableName>>("temperature");
      var_params.set<Real>("gap_conductivity") = getParam<Real>("gap_conductivity");

      if (isParamValid("gap_conductivity_function"))
        var_params.set<FunctionName>("gap_conductivity_function") =
            getParam<FunctionName>("gap_conductivity_function");

      if (isParamValid("gap_conductivity_function_variable"))
        var_params.set<std::vector<VariableName>>("gap_conductivity_function_variable") =
            getParam<std::vector<VariableName>>("gap_conductivity_function_variable");

      var_params.set<Real>("min_gap") = getParam<Real>("min_gap");
      var_params.set<unsigned int>("min_gap_order") = getParam<unsigned int>("min_gap_order");

      var_params.set<std::vector<BoundaryName>>("boundary") =
          getParam<std::vector<BoundaryName>>("boundary");

      var_params.set<bool>("use_displaced_mesh") = true;

      _problem->addUserObject("GapFluxModelConduction",
                              "gap_flux_model_conduction_object_" + MooseUtils::shortName(name()) +
                                  "_" +
                                  Moose::stringify(thermal_action_userobject_conduction_counter++),
                              var_params);
    }
    else if (user_object_vector[i] == "GapFluxModelRadiation")
    {
      auto var_params = _factory.getValidParams("GapFluxModelRadiation");

      var_params.set<Real>("stefan_boltzmann") = getParam<Real>("stefan_boltzmann");
      var_params.set<Real>("primary_emissivity") = getParam<Real>("primary_emissivity");
      var_params.set<Real>("secondary_emissivity") = getParam<Real>("secondary_emissivity");

      var_params.set<std::vector<BoundaryName>>("boundary") =
          getParam<std::vector<BoundaryName>>("boundary");

      var_params.set<std::vector<VariableName>>("temperature") =
          getParam<std::vector<VariableName>>("temperature");

      var_params.set<bool>("use_displaced_mesh") = true;

      _problem->addUserObject("GapFluxModelRadiation",
                              "gap_flux_model_radiation_object_" + MooseUtils::shortName(name()) +
                                  "_" +
                                  Moose::stringify(thermal_action_userobject_radiation_counter++),
                              var_params);
    }
    else
      paramError("user_object_physics",
                 "At least one of the user objects provided to capture gap heat transfer physics "
                 "using a mortar formulation are not currently supported");
  }
}

void
MortarGapHeatTransferAction::checkForExistingSubdomains()
{
  if (parameters().isParamSetByUser("primary_subdomain") &&
      parameters().isParamSetByUser("secondary_subdomain"))
  {
    _user_provided_mortar_meshes = true;
  }
  else
    _user_provided_mortar_meshes = false;
}
<|endoftext|>


<commit_before>#include <boost/format.hpp>
#include <boost/make_shared.hpp>
#include <boost/algorithm/string/predicate.hpp>
#include <boost/range/adaptor/map.hpp>
#include "ManipulatorMarker.h"
#include "or_conversions.h"

using boost::format;
using boost::str;
using boost::adaptors::map_values;
using interactive_markers::InteractiveMarkerServer;
using visualization_msgs::InteractiveMarkerControl;
using visualization_msgs::InteractiveMarkerFeedback;
using visualization_msgs::InteractiveMarkerFeedbackConstPtr;
using OpenRAVE::EnvironmentBasePtr;
using OpenRAVE::RobotBasePtr;
using OpenRAVE::KinBody;
using OpenRAVE::IkSolverBasePtr;

typedef OpenRAVE::RobotBase::RobotStateSaver RobotStateSaver;

typedef boost::shared_ptr<InteractiveMarkerServer> InteractiveMarkerServerPtr;
typedef OpenRAVE::RobotBase::ManipulatorPtr ManipulatorPtr;
typedef OpenRAVE::KinBody::Link::GeometryPtr GeometryPtr;
typedef OpenRAVE::KinBody::LinkPtr LinkPtr;
typedef OpenRAVE::KinBody::LinkInfo LinkInfo;
typedef OpenRAVE::KinBody::LinkInfoConstPtr LinkInfoConstPtr;
typedef OpenRAVE::KinBody::JointPtr JointPtr;
typedef OpenRAVE::KinBody::JointInfo JointInfo;
typedef OpenRAVE::KinBody::JointInfoConstPtr JointInfoConstPtr;
typedef OpenRAVE::RobotBase::ManipulatorInfo ManipulatorInfo;
typedef OpenRAVE::RobotBase::ManipulatorInfoConstPtr ManipulatorInfoConstPtr;
typedef OpenRAVE::RobotBase::AttachedSensorInfo AttachedSensorInfo;
typedef OpenRAVE::RobotBase::AttachedSensorInfoConstPtr AttachedSensorInfoConstPtr;

// TODO: Don't hardcode this.
static std::string const kWorldFrameId = "/world";
static std::string const kGhostKey = "interactive_marker::ghost";

namespace or_interactivemarker {

OpenRAVE::Vector const ManipulatorMarker::kValidColor(0, 1, 0, 0.4);
OpenRAVE::Vector const ManipulatorMarker::kInvalidColor(1, 0, 0, 0.4);

ManipulatorMarker::ManipulatorMarker(InteractiveMarkerServerPtr server,
                                     ManipulatorPtr manipulator)
    : server_(server)
    , manipulator_(manipulator)
    , changed_pose_(true)
    , has_ik_(true)
    , current_pose_(manipulator->GetEndEffectorTransform())
{
    BOOST_ASSERT(server_);
    BOOST_ASSERT(manipulator);

    // Create the ghost manipulator.
    CreateGeometry();

    // Create a 6-DOF pose control.
    ik_marker_.header.frame_id = kWorldFrameId;
    ik_marker_.name = id();
    ik_marker_.description = manipulator_->GetName();
    ik_marker_.pose = toROSPose(manipulator->GetEndEffectorTransform());
    ik_marker_.scale = 0.25;

    {
        InteractiveMarkerControl control;

        control.orientation.w = 1;
        control.orientation.x = 1;
        control.orientation.y = 0;
        control.orientation.z = 0;
        control.name = "rotate_x";
        control.interaction_mode = InteractiveMarkerControl::ROTATE_AXIS;
        ik_marker_.controls.push_back(control);
        control.name = "move_x";
        control.interaction_mode = InteractiveMarkerControl::MOVE_AXIS;
        ik_marker_.controls.push_back(control);

        control.orientation.w = 1;
        control.orientation.x = 0;
        control.orientation.y = 1;
        control.orientation.z = 0;
        control.name = "rotate_y";
        control.interaction_mode = InteractiveMarkerControl::ROTATE_AXIS;
        ik_marker_.controls.push_back(control);
        control.name = "move_y";
        control.interaction_mode = InteractiveMarkerControl::MOVE_AXIS;
        ik_marker_.controls.push_back(control);

        control.orientation.w = 1;
        control.orientation.x = 0;
        control.orientation.y = 0;
        control.orientation.z = 1;
        control.name = "rotate_z";
        control.interaction_mode = InteractiveMarkerControl::ROTATE_AXIS;
        ik_marker_.controls.push_back(control);
        control.name = "move_z";
        control.interaction_mode = InteractiveMarkerControl::MOVE_AXIS;
        ik_marker_.controls.push_back(control);
    }

    server_->insert(ik_marker_);
    server_->setCallback(ik_marker_.name,
        boost::bind(&ManipulatorMarker::IkFeedback, this, _1));

    CreateMenu();
}

ManipulatorMarker::~ManipulatorMarker()
{
    server_->erase(ik_marker_.name);
}

std::string ManipulatorMarker::id() const
{
    OpenRAVE::RobotBasePtr const robot = manipulator_->GetRobot();
    OpenRAVE::EnvironmentBasePtr const env = robot->GetEnv();
    int const environment_id = OpenRAVE::RaveGetEnvironmentId(env);

    return str(format("Environment[%d].KinBody[%s].Manipulator[%s]")
               % environment_id % robot->GetName() % manipulator_->GetName());
}

bool ManipulatorMarker::EnvironmentSync()
{
    ManipulatorPtr const manipulator = manipulator_;
    RobotBasePtr const robot = manipulator->GetRobot();
    IkSolverBasePtr const ik_solver = manipulator->GetIkSolver();
    RobotStateSaver const saver(robot, KinBody::Save_LinkTransformation);

    // Figure out what the free joints are.
    size_t const num_free = ik_solver->GetNumFreeParameters();
    std::vector<JointPtr> free_joints;
    InferFreeJoints(&free_joints);
    BOOST_ASSERT(free_joints.size() == num_free);

    std::vector<int> free_dof_indices;
    for (JointPtr const &free_joint : free_joints) {
        free_dof_indices.push_back(free_joint->GetJointIndex());
    }

    if (ik_solver) {
        RobotStateSaver const free_saver(robot, KinBody::Save_LinkTransformation);

        // Default to the current configuration of the robot.
        if (current_free_.size() != num_free) {
            robot->GetDOFValues(current_free_, free_dof_indices);
        }

        // Extract free parameters from the joint controls.
        bool changed_free = false;
        for (size_t ifree = 0; ifree < num_free; ++ifree) {
            JointPtr const &joint = free_joints[ifree];
            auto const it = free_joint_markers_.find(joint.get());
            if (it == free_joint_markers_.end()) {
                continue;
            }

            // Invalidate the IK solution if we change the free joint.
            JointMarkerPtr &joint_marker = it->second;
            if (joint_marker->angle() != current_free_[ifree]) {
                current_free_[ifree] = joint_marker->angle();
                changed_free = true;
            }
        }

        // Set and clamp these joint values to be within limits.
        robot->SetDOFValues(current_free_, KinBody::CLA_CheckLimitsSilent, free_dof_indices);
        robot->GetDOFValues(current_free_, free_dof_indices);

        // Update our IK solution.
        if (changed_pose_ || changed_free) {
            std::vector<OpenRAVE::dReal> new_ik;
            OpenRAVE::IkParameterization ik_param;
            ik_param.SetTransform6D(current_pose_);

            std::vector<std::vector<OpenRAVE::dReal> > ik_solutions;
            has_ik_ = manipulator_->FindIKSolution(ik_param, new_ik, 0);
            if (has_ik_) {
                current_ik_ = new_ik;
            }
        }
    }
    changed_pose_ = false;

    // Update the pose of the ghost manipulator.
    auto const dof_indices = manipulator->GetArmIndices();
    robot->SetDOFValues(current_ik_, 1, dof_indices);

    bool is_changed = false;
    for (LinkMarkerPtr const &link_marker : link_markers_ | map_values) {
        bool const is_link_changed = link_marker->EnvironmentSync();
        if (!is_link_changed) {
            OpenRAVE::Transform const link_pose = link_marker->link()->GetTransform();
            link_marker->set_pose(link_pose);
        } else {
            UpdateMenu(link_marker);
        }

        // Set the color to indicate whether we have a valid IK solution.
        if (has_ik_) {
            link_marker->set_color(kValidColor);
        } else {
            link_marker->set_color(kInvalidColor);
        }

        is_changed = is_changed || is_link_changed;
    }

    if (ik_solver) {
        for (JointPtr const &joint : free_joints) {
            if (!joint) {
                continue;
            }

            // Lazily create the marker if it is missing.
            JointMarkerPtr &joint_marker = free_joint_markers_[joint.get()];
            if (!joint_marker) {
                joint_marker = boost::make_shared<JointMarker>(server_, joint);
            }

            // Update the pose of the control to match the ghost arm.
            OpenRAVE::Transform const joint_pose = JointMarker::GetJointPose(joint);
            joint_marker->set_pose(joint_pose);

            bool const is_joint_changed = joint_marker->EnvironmentSync();
            is_changed = is_changed || is_joint_changed;
        }
    }
    return is_changed;
}

void ManipulatorMarker::CreateGeometry()
{
    link_markers_.clear();

    ManipulatorPtr const manipulator = manipulator_;
    RobotBasePtr const robot = manipulator->GetRobot();

    // Get links in the manipulator chain.
    std::vector<LinkPtr> chain_links;
    LinkPtr const base_link = manipulator->GetBase();
    LinkPtr const tip_link = manipulator->GetEndEffector();
    bool const success = robot->GetChain(
            base_link->GetIndex(), tip_link->GetIndex(), chain_links);
    BOOST_ASSERT(success);

    // Get end-effector links.
    std::vector<LinkPtr> child_links;
    manipulator->GetChildLinks(child_links);

    // Remove duplicates in case the end-effector link is double-counted.
    std::set<LinkPtr> links;
    links.insert(chain_links.begin(), chain_links.end());
    links.insert(child_links.begin(), child_links.end());

    // Render each link using a LinkMarker.
    for (LinkPtr const &link : links) {
        link_markers_[link.get()] = boost::make_shared<LinkMarker>(server_, link, true);
    }
}

void ManipulatorMarker::CreateMenu()
{
    auto const cb = boost::bind(&ManipulatorMarker::MenuCallback, this, _1);

    menu_set_ = menu_handler_.insert("Set DOF Values", cb);
    menu_reset_ = menu_handler_.insert("Restore DOF Values", cb);
}

void ManipulatorMarker::UpdateMenu()
{
    for (LinkMarkerPtr const &link_marker : link_markers_ | map_values) {
        UpdateMenu(link_marker);
    }
}

void ManipulatorMarker::UpdateMenu(LinkMarkerPtr link_marker)
{
    menu_handler_.apply(*server_, link_marker->interactive_marker()->name);
}

void ManipulatorMarker::MenuCallback(InteractiveMarkerFeedbackConstPtr const &feedback)
{
    ManipulatorPtr const manipulator = manipulator_;
    RobotBasePtr const robot = manipulator->GetRobot();
    std::vector<int> const &arm_indices = manipulator->GetArmIndices();

    if (feedback->menu_entry_id == menu_set_) {
        robot->SetDOFValues(current_ik_, KinBody::CLA_CheckLimits, arm_indices);
        RAVELOG_DEBUG("Set manipulator '%s' to IK solution.\n",
            manipulator->GetName().c_str()
        );
    } else if (feedback->menu_entry_id == menu_reset_) {
        robot->GetDOFValues(current_ik_, arm_indices);
        current_pose_ = manipulator_->GetEndEffectorTransform();
        current_free_.clear();
        RAVELOG_DEBUG("Snapped to current configuration of manipulator '%s'.\n",
            manipulator->GetName().c_str()
        );
    }
}

void ManipulatorMarker::IkFeedback(InteractiveMarkerFeedbackConstPtr const &feedback)
{
    if (feedback->event_type == InteractiveMarkerFeedback::POSE_UPDATE) {
        current_pose_ = toORPose<OpenRAVE::dReal>(feedback->pose);
        changed_pose_ = true;
    }
}

void ManipulatorMarker::InferFreeJoints(std::vector<JointPtr> *free_joints) const
{
    static OpenRAVE::dReal const kEpsilon = 1e-3;

    ManipulatorPtr const manipulator = manipulator_;
    RobotBasePtr const robot = manipulator->GetRobot();
    IkSolverBasePtr const ik_solver = manipulator->GetIkSolver();

    std::vector<OpenRAVE::dReal> lower_limits, upper_limits;
    robot->GetDOFLimits(lower_limits, upper_limits);

    // We'll only accept a joint if it changes the value by at least kEpsilon.
    size_t const num_free = ik_solver->GetNumFreeParameters();
    std::vector<OpenRAVE::dReal> best_ratio(num_free, kEpsilon);
    std::vector<JointPtr> &best_joints = *free_joints;
    best_joints.assign(num_free, JointPtr());

    // Calculate the sensitivity of the free parameters to each joint.
    std::vector<int> const &arm_indices = manipulator->GetArmIndices();
    std::vector<OpenRAVE::dReal> free_motion;
    for (int const dof_index : arm_indices) {
        RobotStateSaver const saver(robot, KinBody::Save_LinkTransformation);
        std::vector<OpenRAVE::dReal> dof_values;
        robot->GetDOFValues(dof_values);

        // Move the joint to its limits.
        std::vector<OpenRAVE::dReal> lower_free;
        dof_values[dof_index] = lower_limits[dof_index];
        robot->SetDOFValues(dof_values);
        ik_solver->GetFreeParameters(lower_free);
        BOOST_ASSERT(lower_free.size() == num_free);

        std::vector<OpenRAVE::dReal> upper_free;
        dof_values[dof_index] = upper_limits[dof_index];
        robot->SetDOFValues(dof_values);
        ik_solver->GetFreeParameters(upper_free);
        BOOST_ASSERT(upper_free.size() == num_free);

        // Calculate the relative change in free parameters.
        for (size_t ifree = 0; ifree < num_free; ++ifree) {
            double const delta_param = upper_free[ifree] - lower_free[ifree];
            double const delta_value = upper_limits[ifree] - lower_limits[ifree];
            BOOST_ASSERT(delta_value > 0);
            double const ratio = std::fabs(delta_param / delta_value);

            if (ratio > best_ratio[ifree]) {
                best_ratio[ifree] = ratio;
                best_joints[ifree] = robot->GetJointFromDOFIndex(dof_index);
            }
        }
    }
}

}
<commit_msg>Hack to avoid crashing with a NULL IK solver (#3)<commit_after>#include <boost/format.hpp>
#include <boost/make_shared.hpp>
#include <boost/algorithm/string/predicate.hpp>
#include <boost/range/adaptor/map.hpp>
#include "ManipulatorMarker.h"
#include "or_conversions.h"

using boost::format;
using boost::str;
using boost::adaptors::map_values;
using interactive_markers::InteractiveMarkerServer;
using visualization_msgs::InteractiveMarkerControl;
using visualization_msgs::InteractiveMarkerFeedback;
using visualization_msgs::InteractiveMarkerFeedbackConstPtr;
using OpenRAVE::EnvironmentBasePtr;
using OpenRAVE::RobotBasePtr;
using OpenRAVE::KinBody;
using OpenRAVE::IkSolverBasePtr;

typedef OpenRAVE::RobotBase::RobotStateSaver RobotStateSaver;

typedef boost::shared_ptr<InteractiveMarkerServer> InteractiveMarkerServerPtr;
typedef OpenRAVE::RobotBase::ManipulatorPtr ManipulatorPtr;
typedef OpenRAVE::KinBody::Link::GeometryPtr GeometryPtr;
typedef OpenRAVE::KinBody::LinkPtr LinkPtr;
typedef OpenRAVE::KinBody::LinkInfo LinkInfo;
typedef OpenRAVE::KinBody::LinkInfoConstPtr LinkInfoConstPtr;
typedef OpenRAVE::KinBody::JointPtr JointPtr;
typedef OpenRAVE::KinBody::JointInfo JointInfo;
typedef OpenRAVE::KinBody::JointInfoConstPtr JointInfoConstPtr;
typedef OpenRAVE::RobotBase::ManipulatorInfo ManipulatorInfo;
typedef OpenRAVE::RobotBase::ManipulatorInfoConstPtr ManipulatorInfoConstPtr;
typedef OpenRAVE::RobotBase::AttachedSensorInfo AttachedSensorInfo;
typedef OpenRAVE::RobotBase::AttachedSensorInfoConstPtr AttachedSensorInfoConstPtr;

// TODO: Don't hardcode this.
static std::string const kWorldFrameId = "/world";
static std::string const kGhostKey = "interactive_marker::ghost";

namespace or_interactivemarker {

OpenRAVE::Vector const ManipulatorMarker::kValidColor(0, 1, 0, 0.4);
OpenRAVE::Vector const ManipulatorMarker::kInvalidColor(1, 0, 0, 0.4);

ManipulatorMarker::ManipulatorMarker(InteractiveMarkerServerPtr server,
                                     ManipulatorPtr manipulator)
    : server_(server)
    , manipulator_(manipulator)
    , changed_pose_(true)
    , has_ik_(true)
    , current_pose_(manipulator->GetEndEffectorTransform())
{
    BOOST_ASSERT(server_);
    BOOST_ASSERT(manipulator);

    // Create the ghost manipulator.
    CreateGeometry();

    // Create a 6-DOF pose control.
    ik_marker_.header.frame_id = kWorldFrameId;
    ik_marker_.name = id();
    ik_marker_.description = manipulator_->GetName();
    ik_marker_.pose = toROSPose(manipulator->GetEndEffectorTransform());
    ik_marker_.scale = 0.25;

    {
        InteractiveMarkerControl control;

        control.orientation.w = 1;
        control.orientation.x = 1;
        control.orientation.y = 0;
        control.orientation.z = 0;
        control.name = "rotate_x";
        control.interaction_mode = InteractiveMarkerControl::ROTATE_AXIS;
        ik_marker_.controls.push_back(control);
        control.name = "move_x";
        control.interaction_mode = InteractiveMarkerControl::MOVE_AXIS;
        ik_marker_.controls.push_back(control);

        control.orientation.w = 1;
        control.orientation.x = 0;
        control.orientation.y = 1;
        control.orientation.z = 0;
        control.name = "rotate_y";
        control.interaction_mode = InteractiveMarkerControl::ROTATE_AXIS;
        ik_marker_.controls.push_back(control);
        control.name = "move_y";
        control.interaction_mode = InteractiveMarkerControl::MOVE_AXIS;
        ik_marker_.controls.push_back(control);

        control.orientation.w = 1;
        control.orientation.x = 0;
        control.orientation.y = 0;
        control.orientation.z = 1;
        control.name = "rotate_z";
        control.interaction_mode = InteractiveMarkerControl::ROTATE_AXIS;
        ik_marker_.controls.push_back(control);
        control.name = "move_z";
        control.interaction_mode = InteractiveMarkerControl::MOVE_AXIS;
        ik_marker_.controls.push_back(control);
    }

    server_->insert(ik_marker_);
    server_->setCallback(ik_marker_.name,
        boost::bind(&ManipulatorMarker::IkFeedback, this, _1));

    CreateMenu();
}

ManipulatorMarker::~ManipulatorMarker()
{
    server_->erase(ik_marker_.name);
}

std::string ManipulatorMarker::id() const
{
    OpenRAVE::RobotBasePtr const robot = manipulator_->GetRobot();
    OpenRAVE::EnvironmentBasePtr const env = robot->GetEnv();
    int const environment_id = OpenRAVE::RaveGetEnvironmentId(env);

    return str(format("Environment[%d].KinBody[%s].Manipulator[%s]")
               % environment_id % robot->GetName() % manipulator_->GetName());
}

bool ManipulatorMarker::EnvironmentSync()
{
    ManipulatorPtr const manipulator = manipulator_;
    RobotBasePtr const robot = manipulator->GetRobot();
    IkSolverBasePtr const ik_solver = manipulator->GetIkSolver();
    RobotStateSaver const saver(robot, KinBody::Save_LinkTransformation);

    // Hack to avoid crashing if no IK solver is set.
    if (!ik_solver) {
        return false;
    }

    // Figure out what the free joints are.
    size_t const num_free = ik_solver->GetNumFreeParameters();
    std::vector<JointPtr> free_joints;
    InferFreeJoints(&free_joints);
    BOOST_ASSERT(free_joints.size() == num_free);

    std::vector<int> free_dof_indices;
    for (JointPtr const &free_joint : free_joints) {
        free_dof_indices.push_back(free_joint->GetJointIndex());
    }

    if (ik_solver) {
        RobotStateSaver const free_saver(robot, KinBody::Save_LinkTransformation);

        // Default to the current configuration of the robot.
        if (current_free_.size() != num_free) {
            robot->GetDOFValues(current_free_, free_dof_indices);
        }

        // Extract free parameters from the joint controls.
        bool changed_free = false;
        for (size_t ifree = 0; ifree < num_free; ++ifree) {
            JointPtr const &joint = free_joints[ifree];
            auto const it = free_joint_markers_.find(joint.get());
            if (it == free_joint_markers_.end()) {
                continue;
            }

            // Invalidate the IK solution if we change the free joint.
            JointMarkerPtr &joint_marker = it->second;
            if (joint_marker->angle() != current_free_[ifree]) {
                current_free_[ifree] = joint_marker->angle();
                changed_free = true;
            }
        }

        // Set and clamp these joint values to be within limits.
        robot->SetDOFValues(current_free_, KinBody::CLA_CheckLimitsSilent, free_dof_indices);
        robot->GetDOFValues(current_free_, free_dof_indices);

        // Update our IK solution.
        if (changed_pose_ || changed_free) {
            std::vector<OpenRAVE::dReal> new_ik;
            OpenRAVE::IkParameterization ik_param;
            ik_param.SetTransform6D(current_pose_);

            std::vector<std::vector<OpenRAVE::dReal> > ik_solutions;
            has_ik_ = manipulator_->FindIKSolution(ik_param, new_ik, 0);
            if (has_ik_) {
                current_ik_ = new_ik;
            }
        }
    }
    changed_pose_ = false;

    // Update the pose of the ghost manipulator.
    auto const dof_indices = manipulator->GetArmIndices();
    robot->SetDOFValues(current_ik_, 1, dof_indices);

    bool is_changed = false;
    for (LinkMarkerPtr const &link_marker : link_markers_ | map_values) {
        bool const is_link_changed = link_marker->EnvironmentSync();
        if (!is_link_changed) {
            OpenRAVE::Transform const link_pose = link_marker->link()->GetTransform();
            link_marker->set_pose(link_pose);
        } else {
            UpdateMenu(link_marker);
        }

        // Set the color to indicate whether we have a valid IK solution.
        if (has_ik_) {
            link_marker->set_color(kValidColor);
        } else {
            link_marker->set_color(kInvalidColor);
        }

        is_changed = is_changed || is_link_changed;
    }

    if (ik_solver) {
        for (JointPtr const &joint : free_joints) {
            if (!joint) {
                continue;
            }

            // Lazily create the marker if it is missing.
            JointMarkerPtr &joint_marker = free_joint_markers_[joint.get()];
            if (!joint_marker) {
                joint_marker = boost::make_shared<JointMarker>(server_, joint);
            }

            // Update the pose of the control to match the ghost arm.
            OpenRAVE::Transform const joint_pose = JointMarker::GetJointPose(joint);
            joint_marker->set_pose(joint_pose);

            bool const is_joint_changed = joint_marker->EnvironmentSync();
            is_changed = is_changed || is_joint_changed;
        }
    }
    return is_changed;
}

void ManipulatorMarker::CreateGeometry()
{
    link_markers_.clear();

    ManipulatorPtr const manipulator = manipulator_;
    RobotBasePtr const robot = manipulator->GetRobot();

    // Get links in the manipulator chain.
    std::vector<LinkPtr> chain_links;
    LinkPtr const base_link = manipulator->GetBase();
    LinkPtr const tip_link = manipulator->GetEndEffector();
    bool const success = robot->GetChain(
            base_link->GetIndex(), tip_link->GetIndex(), chain_links);
    BOOST_ASSERT(success);

    // Get end-effector links.
    std::vector<LinkPtr> child_links;
    manipulator->GetChildLinks(child_links);

    // Remove duplicates in case the end-effector link is double-counted.
    std::set<LinkPtr> links;
    links.insert(chain_links.begin(), chain_links.end());
    links.insert(child_links.begin(), child_links.end());

    // Render each link using a LinkMarker.
    for (LinkPtr const &link : links) {
        link_markers_[link.get()] = boost::make_shared<LinkMarker>(server_, link, true);
    }
}

void ManipulatorMarker::CreateMenu()
{
    auto const cb = boost::bind(&ManipulatorMarker::MenuCallback, this, _1);

    menu_set_ = menu_handler_.insert("Set DOF Values", cb);
    menu_reset_ = menu_handler_.insert("Restore DOF Values", cb);
}

void ManipulatorMarker::UpdateMenu()
{
    for (LinkMarkerPtr const &link_marker : link_markers_ | map_values) {
        UpdateMenu(link_marker);
    }
}

void ManipulatorMarker::UpdateMenu(LinkMarkerPtr link_marker)
{
    menu_handler_.apply(*server_, link_marker->interactive_marker()->name);
}

void ManipulatorMarker::MenuCallback(InteractiveMarkerFeedbackConstPtr const &feedback)
{
    ManipulatorPtr const manipulator = manipulator_;
    RobotBasePtr const robot = manipulator->GetRobot();
    std::vector<int> const &arm_indices = manipulator->GetArmIndices();

    if (feedback->menu_entry_id == menu_set_) {
        robot->SetDOFValues(current_ik_, KinBody::CLA_CheckLimits, arm_indices);
        RAVELOG_DEBUG("Set manipulator '%s' to IK solution.\n",
            manipulator->GetName().c_str()
        );
    } else if (feedback->menu_entry_id == menu_reset_) {
        robot->GetDOFValues(current_ik_, arm_indices);
        current_pose_ = manipulator_->GetEndEffectorTransform();
        current_free_.clear();
        RAVELOG_DEBUG("Snapped to current configuration of manipulator '%s'.\n",
            manipulator->GetName().c_str()
        );
    }
}

void ManipulatorMarker::IkFeedback(InteractiveMarkerFeedbackConstPtr const &feedback)
{
    if (feedback->event_type == InteractiveMarkerFeedback::POSE_UPDATE) {
        current_pose_ = toORPose<OpenRAVE::dReal>(feedback->pose);
        changed_pose_ = true;
    }
}

void ManipulatorMarker::InferFreeJoints(std::vector<JointPtr> *free_joints) const
{
    static OpenRAVE::dReal const kEpsilon = 1e-3;

    ManipulatorPtr const manipulator = manipulator_;
    RobotBasePtr const robot = manipulator->GetRobot();
    IkSolverBasePtr const ik_solver = manipulator->GetIkSolver();

    std::vector<OpenRAVE::dReal> lower_limits, upper_limits;
    robot->GetDOFLimits(lower_limits, upper_limits);

    // We'll only accept a joint if it changes the value by at least kEpsilon.
    size_t const num_free = ik_solver->GetNumFreeParameters();
    std::vector<OpenRAVE::dReal> best_ratio(num_free, kEpsilon);
    std::vector<JointPtr> &best_joints = *free_joints;
    best_joints.assign(num_free, JointPtr());

    // Calculate the sensitivity of the free parameters to each joint.
    std::vector<int> const &arm_indices = manipulator->GetArmIndices();
    std::vector<OpenRAVE::dReal> free_motion;
    for (int const dof_index : arm_indices) {
        RobotStateSaver const saver(robot, KinBody::Save_LinkTransformation);
        std::vector<OpenRAVE::dReal> dof_values;
        robot->GetDOFValues(dof_values);

        // Move the joint to its limits.
        std::vector<OpenRAVE::dReal> lower_free;
        dof_values[dof_index] = lower_limits[dof_index];
        robot->SetDOFValues(dof_values);
        ik_solver->GetFreeParameters(lower_free);
        BOOST_ASSERT(lower_free.size() == num_free);

        std::vector<OpenRAVE::dReal> upper_free;
        dof_values[dof_index] = upper_limits[dof_index];
        robot->SetDOFValues(dof_values);
        ik_solver->GetFreeParameters(upper_free);
        BOOST_ASSERT(upper_free.size() == num_free);

        // Calculate the relative change in free parameters.
        for (size_t ifree = 0; ifree < num_free; ++ifree) {
            double const delta_param = upper_free[ifree] - lower_free[ifree];
            double const delta_value = upper_limits[ifree] - lower_limits[ifree];
            BOOST_ASSERT(delta_value > 0);
            double const ratio = std::fabs(delta_param / delta_value);

            if (ratio > best_ratio[ifree]) {
                best_ratio[ifree] = ratio;
                best_joints[ifree] = robot->GetJointFromDOFIndex(dof_index);
            }
        }
    }
}

}
<|endoftext|>