├── .gitignore ├── CMakeLists.txt ├── README.md ├── automate ├── __init__.py ├── arg_parsing.py ├── brep.py ├── conversions.py ├── eclasses.py ├── plot_confusion_matrix.py ├── pointnet_encoder.py ├── sbgcn.py ├── util.py └── uvnet_encoders.py ├── cpp ├── automate.cpp ├── disjointset.cpp ├── disjointset.h ├── eclass.cpp ├── eclass.h ├── lsh.cpp ├── lsh.h ├── part.cpp └── part.h ├── environment.yml ├── minimal_env.yml └── setup.py /.gitignore: -------------------------------------------------------------------------------- 1 | checkpoints/ 2 | *.swp 3 | .vscode/ 4 | .vs/ 5 | .DS_Store 6 | .idea/ 7 | 8 | # Byte-compiled / optimized / DLL files 9 | __pycache__/ 10 | *.py[cod] 11 | *$py.class 12 | 13 | # Data and Notebook scratch directories 14 | data/* 15 | notebooks/* 16 | 17 | # Pytorch Lightning 18 | *lightning_logs* 19 | *experiments* 20 | 21 | # C extensions 22 | *.so 23 | 24 | # Distribution / packaging 25 | .Python 26 | build/ 27 | develop-eggs/ 28 | dist/ 29 | downloads/ 30 | eggs/ 31 | .eggs/ 32 | lib/ 33 | lib64/ 34 | parts/ 35 | sdist/ 36 | var/ 37 | wheels/ 38 | share/python-wheels/ 39 | *.egg-info/ 40 | .installed.cfg 41 | *.egg 42 | MANIFEST 43 | 44 | # PyInstaller 45 | # Usually these files are written by a python script from a template 46 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 47 | *.manifest 48 | *.spec 49 | 50 | # Installer logs 51 | pip-log.txt 52 | pip-delete-this-directory.txt 53 | 54 | # Unit test / coverage reports 55 | htmlcov/ 56 | .tox/ 57 | .nox/ 58 | .coverage 59 | .coverage.* 60 | .cache 61 | nosetests.xml 62 | coverage.xml 63 | *.cover 64 | *.py,cover 65 | .hypothesis/ 66 | .pytest_cache/ 67 | cover/ 68 | 69 | # Translations 70 | *.mo 71 | *.pot 72 | 73 | # Django stuff: 74 | *.log 75 | local_settings.py 76 | db.sqlite3 77 | db.sqlite3-journal 78 | 79 | # Flask stuff: 80 | instance/ 81 | .webassets-cache 82 | 83 | # Scrapy stuff: 84 | .scrapy 85 | 86 | # Sphinx documentation 87 | docs/_build/ 88 | 89 | # PyBuilder 90 | .pybuilder/ 91 | target/ 92 | 93 | # Jupyter Notebook 94 | .ipynb_checkpoints 95 | *.ipynb 96 | 97 | # IPython 98 | profile_default/ 99 | ipython_config.py 100 | 101 | # pyenv 102 | # For a library or package, you might want to ignore these files since the code is 103 | # intended to run in multiple environments; otherwise, check them in: 104 | # .python-version 105 | 106 | # pipenv 107 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 108 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 109 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 110 | # install all needed dependencies. 111 | #Pipfile.lock 112 | 113 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 114 | __pypackages__/ 115 | 116 | # Celery stuff 117 | celerybeat-schedule 118 | celerybeat.pid 119 | 120 | # SageMath parsed files 121 | *.sage.py 122 | 123 | # Environments 124 | .env 125 | .venv 126 | env/ 127 | venv/ 128 | ENV/ 129 | env.bak/ 130 | venv.bak/ 131 | 132 | # Spyder project settings 133 | .spyderproject 134 | .spyproject 135 | 136 | # Rope project settings 137 | .ropeproject 138 | 139 | # mkdocs documentation 140 | /site 141 | 142 | # mypy 143 | .mypy_cache/ 144 | .dmypy.json 145 | dmypy.json 146 | 147 | # Pyre type checker 148 | .pyre/ 149 | 150 | # pytype static type analyzer 151 | .pytype/ 152 | 153 | # Cython debug symbols 154 | cython_debug/ 155 | 156 | *.bak 157 | 158 | # Visual Studio 159 | enc_temp_folder/ 160 | -------------------------------------------------------------------------------- /CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.14) 2 | project(automate_cpp) 3 | 4 | set(CMAKE_CXX_STANDARD 14) 5 | 6 | set(CMAKE_POSITION_INDEPENDENT_CODE ON) 7 | 8 | find_package(pybind11 CONFIG REQUIRED) 9 | find_package(Eigen3 REQUIRED) 10 | 11 | include(FetchContent) 12 | FetchContent_Declare( 13 | brloader 14 | GIT_REPOSITORY https://github.com/deGravity/breploader.git 15 | GIT_TAG v0.5 16 | ) 17 | 18 | FetchContent_MakeAvailable(brloader) 19 | 20 | pybind11_add_module(automate_cpp 21 | ${CMAKE_CURRENT_SOURCE_DIR}/cpp/automate.cpp 22 | ${CMAKE_CURRENT_SOURCE_DIR}/cpp/disjointset.h 23 | ${CMAKE_CURRENT_SOURCE_DIR}/cpp/disjointset.cpp 24 | ${CMAKE_CURRENT_SOURCE_DIR}/cpp/lsh.h 25 | ${CMAKE_CURRENT_SOURCE_DIR}/cpp/lsh.cpp 26 | ${CMAKE_CURRENT_SOURCE_DIR}/cpp/eclass.h 27 | ${CMAKE_CURRENT_SOURCE_DIR}/cpp/eclass.cpp 28 | ${CMAKE_CURRENT_SOURCE_DIR}/cpp/part.h 29 | ${CMAKE_CURRENT_SOURCE_DIR}/cpp/part.cpp 30 | ) 31 | 32 | target_link_libraries(automate_cpp PUBLIC breploader Eigen3::Eigen) -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # AutoMate 2 | Dataset and code for the automatic mating of CAD assemblies. 3 | 4 | ## Paper 5 | Please see our [paper](https://dl.acm.org/doi/10.1145/3478513.3480562) for more details. 6 | 7 | ## Dataset 8 | The AutoMate Dataset can be downloaded from [Zenodo](https://zenodo.org/record/7776208#.ZDcYinbMIQ8). 9 | 10 | 11 | ## Installation 12 | 13 | Automate relies on C++ extension modules to read Parasolid and STEP files. Since Parasolid is a proprietary CAD kernel, we can't distribute it, so you need to have the distribution on your machine already and compile it at install time. 14 | 15 | ### Requirements 16 | 17 | Installation relies on CMake, OpenCascade, and Parasolid. In the future, we intend to make the Parasolid dependency optional. The easiest way to get the first two dependencies (and all python dependencies) is to install the conda environments `environment.yml` or `minimal_env.yml`: 18 | 19 | `conda env create -f [environment|minimal_env].yml` 20 | 21 | The Parasolid requirement relies on setting the environmental variable `$PARASOLID_BASE` on your system pointing to the Parasolid install directory for your operating system. For example 22 | 23 | ``export PARASOLID_BASE=${PATH_TO_PARASOLID_INSTALL}/intel_linux/base`` 24 | 25 | Replace ``intel_linux`` with the directory appropriate to your OS. The base directory should contain files like `pskernel_archive.lib` and `parasolid_kernel.h`. 26 | 27 | Once these requirements are met, you an install via pip: 28 | 29 | `pip install git+https://github.com/degravity/automate.git@v1.0.4` 30 | 31 | 32 | ### Troubleshooting 33 | 34 | ``` 35 | ImportError: dynamic module does not define module export function (PyInit_automate_cpp) 36 | ``` 37 | 38 | If you get this error when trying to import part of the module, it means that python can't find the C++ extensions module. To fix this, try cleaning out **all** build files and building again. 39 | 40 | ## Citing 41 | 42 | If you use this code our the [AutoMate Dataset](https://zenodo.org/record/7776208#.ZDcYinbMIQ8) in your work, please cite us: 43 | 44 | ``` 45 | @article{10.1145/3478513.3480562, 46 | author = {Jones, Benjamin and Hildreth, Dalton and Chen, Duowen and Baran, Ilya and Kim, Vladimir G. and Schulz, Adriana}, 47 | title = {AutoMate: A Dataset and Learning Approach for Automatic Mating of CAD Assemblies}, 48 | year = {2021}, 49 | issue_date = {December 2021}, 50 | publisher = {Association for Computing Machinery}, 51 | address = {New York, NY, USA}, 52 | volume = {40}, 53 | number = {6}, 54 | issn = {0730-0301}, 55 | url = {https://doi.org/10.1145/3478513.3480562}, 56 | doi = {10.1145/3478513.3480562}, 57 | month = {dec}, 58 | articleno = {227}, 59 | numpages = {18}, 60 | keywords = {assembly-based modeling, representation learning, boundary representation, computer-aided design} 61 | } 62 | ``` 63 | -------------------------------------------------------------------------------- /automate/__init__.py: -------------------------------------------------------------------------------- 1 | from .conversions import jsonify, torchify 2 | from .brep import PartFeatures, part_to_graph, HetData, PartDataset, flatbatch 3 | from .sbgcn import SBGCN, LinearBlock 4 | 5 | from .util import run_model, ArgparseInitialized 6 | from .brep import PartFeatures, part_to_graph, HetData, PartDataset 7 | from .sbgcn import SBGCN, LinearBlock, BipartiteResMRConv 8 | 9 | from .util import run_model, ArgparseInitialized 10 | from .eclasses import find_eclasses 11 | 12 | from automate_cpp import Part, PartOptions 13 | 14 | 15 | __all__ = [ 16 | 'jsonify', 17 | 'torchify', 18 | 'PartFeatures', 19 | 'part_to_graph', 20 | 'HetData', 21 | 'SBGCN', 22 | 'LinearBlock', 23 | 'PartDataset', 24 | 'flatbatch', 25 | 'run_model', 26 | 'ArgparseInitialized', 27 | 'BipartiteResMRConv', 28 | 'Part', 29 | 'PartOptions', 30 | 'find_eclasses' 31 | ] -------------------------------------------------------------------------------- /automate/arg_parsing.py: -------------------------------------------------------------------------------- 1 | # Copied and Modified from PyTorch Lightning 2 | # 3 | # Original Copyright Notice: 4 | # 5 | # ======================================================================== 6 | # Copyright The PyTorch Lightning team. 7 | # 8 | # Licensed under the Apache License, Version 2.0 (the "License"); 9 | # you may not use this file except in compliance with the License. 10 | # You may obtain a copy of the License at 11 | # 12 | # http://www.apache.org/licenses/LICENSE-2.0 13 | # 14 | # Unless required by applicable law or agreed to in writing, software 15 | # distributed under the License is distributed on an "AS IS" BASIS, 16 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 17 | # See the License for the specific language governing permissions and 18 | # limitations under the License. 19 | # ======================================================================== 20 | 21 | 22 | import inspect 23 | import os 24 | from abc import ABC 25 | from argparse import _ArgumentGroup, ArgumentParser, Namespace 26 | from contextlib import suppress 27 | from functools import wraps 28 | from typing import Any, Callable, Dict, List, Tuple, Type, Union 29 | 30 | 31 | 32 | def str_to_bool_or_str(val: str) -> Union[str, bool]: 33 | """Possibly convert a string representation of truth to bool. Returns the input otherwise. Based on the python 34 | implementation distutils.utils.strtobool. 35 | True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values are 'n', 'no', 'f', 'false', 'off', and '0'. 36 | """ 37 | lower = val.lower() 38 | if lower in ("y", "yes", "t", "true", "on", "1"): 39 | return True 40 | if lower in ("n", "no", "f", "false", "off", "0"): 41 | return False 42 | return val 43 | 44 | 45 | def str_to_bool(val: str) -> bool: 46 | """Convert a string representation of truth to bool. 47 | True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values 48 | are 'n', 'no', 'f', 'false', 'off', and '0'. 49 | Raises: 50 | ValueError: 51 | If ``val`` isn't in one of the aforementioned true or false values. 52 | >>> str_to_bool('YES') 53 | True 54 | >>> str_to_bool('FALSE') 55 | False 56 | """ 57 | val_converted = str_to_bool_or_str(val) 58 | if isinstance(val_converted, bool): 59 | return val_converted 60 | raise ValueError(f"invalid truth value {val_converted}") 61 | 62 | 63 | def str_to_bool_or_int(val: str) -> Union[bool, int, str]: 64 | """Convert a string representation to truth of bool if possible, or otherwise try to convert it to an int. 65 | >>> str_to_bool_or_int("FALSE") 66 | False 67 | >>> str_to_bool_or_int("1") 68 | True 69 | >>> str_to_bool_or_int("2") 70 | 2 71 | >>> str_to_bool_or_int("abc") 72 | 'abc' 73 | """ 74 | val_converted = str_to_bool_or_str(val) 75 | if isinstance(val_converted, bool): 76 | return val_converted 77 | try: 78 | return int(val_converted) 79 | except ValueError: 80 | return val_converted 81 | 82 | class ParseArgparserDataType(ABC): 83 | def __init__(self, *_: Any, **__: Any) -> None: 84 | pass 85 | 86 | @classmethod 87 | def parse_argparser(cls, args: "ArgumentParser") -> Any: 88 | pass 89 | 90 | 91 | def from_argparse_args( 92 | cls: Type[ParseArgparserDataType], args: Union[Namespace, ArgumentParser], **kwargs: Any 93 | ) -> ParseArgparserDataType: 94 | """Create an instance from CLI arguments. Eventually use varibles from OS environement which are defined as 95 | ``"PL__"``. 96 | Args: 97 | cls: Lightning class 98 | args: The parser or namespace to take arguments from. Only known arguments will be 99 | parsed and passed to the :class:`Trainer`. 100 | **kwargs: Additional keyword arguments that may override ones in the parser or namespace. 101 | These must be valid Trainer arguments. 102 | Examples: 103 | >>> from pytorch_lightning import Trainer 104 | >>> parser = ArgumentParser(add_help=False) 105 | >>> parser = Trainer.add_argparse_args(parser) 106 | >>> parser.add_argument('--my_custom_arg', default='something') # doctest: +SKIP 107 | >>> args = Trainer.parse_argparser(parser.parse_args("")) 108 | >>> trainer = Trainer.from_argparse_args(args, logger=False) 109 | """ 110 | if isinstance(args, ArgumentParser): 111 | args = cls.parse_argparser(args) 112 | 113 | params = vars(args) 114 | 115 | # we only want to pass in valid Trainer args, the rest may be user specific 116 | valid_kwargs = inspect.signature(cls.__init__).parameters 117 | trainer_kwargs = {name: params[name] for name in valid_kwargs if name in params} 118 | trainer_kwargs.update(**kwargs) 119 | 120 | return cls(**trainer_kwargs) 121 | 122 | 123 | def parse_argparser(cls: Type["pl.Trainer"], arg_parser: Union[ArgumentParser, Namespace]) -> Namespace: 124 | """Parse CLI arguments, required for custom bool types.""" 125 | args = arg_parser.parse_args() if isinstance(arg_parser, ArgumentParser) else arg_parser 126 | 127 | types_default = {arg: (arg_types, arg_default) for arg, arg_types, arg_default in get_init_arguments_and_types(cls)} 128 | 129 | modified_args = {} 130 | for k, v in vars(args).items(): 131 | if k in types_default and v is None: 132 | # We need to figure out if the None is due to using nargs="?" or if it comes from the default value 133 | arg_types, arg_default = types_default[k] 134 | if bool in arg_types and isinstance(arg_default, bool): 135 | # Value has been passed as a flag => It is currently None, so we need to set it to True 136 | # We always set to True, regardless of the default value. 137 | # Users must pass False directly, but when passing nothing True is assumed. 138 | # i.e. the only way to disable something that defaults to True is to use the long form: 139 | # "--a_default_true_arg False" becomes False, while "--a_default_false_arg" becomes None, 140 | # which then becomes True here. 141 | 142 | v = True 143 | 144 | modified_args[k] = v 145 | return Namespace(**modified_args) 146 | 147 | 148 | def parse_env_variables(cls: Type["pl.Trainer"], template: str = "PL_%(cls_name)s_%(cls_argument)s") -> Namespace: 149 | """Parse environment arguments if they are defined. 150 | Examples: 151 | >>> from pytorch_lightning import Trainer 152 | >>> parse_env_variables(Trainer) 153 | Namespace() 154 | >>> import os 155 | >>> os.environ["PL_TRAINER_GPUS"] = '42' 156 | >>> os.environ["PL_TRAINER_BLABLABLA"] = '1.23' 157 | >>> parse_env_variables(Trainer) 158 | Namespace(gpus=42) 159 | >>> del os.environ["PL_TRAINER_GPUS"] 160 | """ 161 | cls_arg_defaults = get_init_arguments_and_types(cls) 162 | 163 | env_args = {} 164 | for arg_name, _, _ in cls_arg_defaults: 165 | env = template % {"cls_name": cls.__name__.upper(), "cls_argument": arg_name.upper()} 166 | val = os.environ.get(env) 167 | if not (val is None or val == ""): 168 | # todo: specify the possible exception 169 | with suppress(Exception): 170 | # converting to native types like int/float/bool 171 | val = eval(val) 172 | env_args[arg_name] = val 173 | return Namespace(**env_args) 174 | 175 | 176 | def get_init_arguments_and_types(cls: Any) -> List[Tuple[str, Tuple, Any]]: 177 | r"""Scans the class signature and returns argument names, types and default values. 178 | Returns: 179 | List with tuples of 3 values: 180 | (argument name, set with argument types, argument default value). 181 | Examples: 182 | >>> from pytorch_lightning import Trainer 183 | >>> args = get_init_arguments_and_types(Trainer) 184 | """ 185 | cls_default_params = inspect.signature(cls).parameters 186 | name_type_default = [] 187 | for arg in cls_default_params: 188 | arg_type = cls_default_params[arg].annotation 189 | arg_default = cls_default_params[arg].default 190 | try: 191 | arg_types = tuple(arg_type.__args__) 192 | except (AttributeError, TypeError): 193 | arg_types = (arg_type,) 194 | 195 | name_type_default.append((arg, arg_types, arg_default)) 196 | 197 | return name_type_default 198 | 199 | 200 | def _get_abbrev_qualified_cls_name(cls: Any) -> str: 201 | assert isinstance(cls, type), repr(cls) 202 | if cls.__module__.startswith("pytorch_lightning."): 203 | # Abbreviate. 204 | return f"pl.{cls.__name__}" 205 | # Fully qualified. 206 | return f"{cls.__module__}.{cls.__qualname__}" 207 | 208 | 209 | def add_argparse_args( 210 | cls: Type["pl.Trainer"], parent_parser: ArgumentParser, *, use_argument_group: bool = True 211 | ) -> Union[_ArgumentGroup, ArgumentParser]: 212 | r"""Extends existing argparse by default attributes for ``cls``. 213 | Args: 214 | cls: Lightning class 215 | parent_parser: 216 | The custom cli arguments parser, which will be extended by 217 | the class's default arguments. 218 | use_argument_group: 219 | By default, this is True, and uses ``add_argument_group`` to add 220 | a new group. 221 | If False, this will use old behavior. 222 | Returns: 223 | If use_argument_group is True, returns ``parent_parser`` to keep old 224 | workflows. If False, will return the new ArgumentParser object. 225 | Only arguments of the allowed types (str, float, int, bool) will 226 | extend the ``parent_parser``. 227 | Raises: 228 | RuntimeError: 229 | If ``parent_parser`` is not an ``ArgumentParser`` instance 230 | Examples: 231 | >>> # Option 1: Default usage. 232 | >>> import argparse 233 | >>> from pytorch_lightning import Trainer 234 | >>> parser = argparse.ArgumentParser() 235 | >>> parser = Trainer.add_argparse_args(parser) 236 | >>> args = parser.parse_args([]) 237 | >>> # Option 2: Disable use_argument_group (old behavior). 238 | >>> import argparse 239 | >>> from pytorch_lightning import Trainer 240 | >>> parser = argparse.ArgumentParser() 241 | >>> parser = Trainer.add_argparse_args(parser, use_argument_group=False) 242 | >>> args = parser.parse_args([]) 243 | """ 244 | if isinstance(parent_parser, _ArgumentGroup): 245 | raise RuntimeError("Please only pass an ArgumentParser instance.") 246 | if use_argument_group: 247 | group_name = _get_abbrev_qualified_cls_name(cls) 248 | parser: Union[_ArgumentGroup, ArgumentParser] = parent_parser.add_argument_group(group_name) 249 | else: 250 | parser = ArgumentParser(parents=[parent_parser], add_help=False) 251 | 252 | ignore_arg_names = ["self", "args", "kwargs"] 253 | if hasattr(cls, "get_deprecated_arg_names"): 254 | ignore_arg_names += cls.get_deprecated_arg_names() 255 | 256 | allowed_types = (str, int, float, bool) 257 | 258 | # Get symbols from cls or init function. 259 | for symbol in (cls, cls.__init__): 260 | args_and_types = get_init_arguments_and_types(symbol) 261 | args_and_types = [x for x in args_and_types if x[0] not in ignore_arg_names] 262 | if len(args_and_types) > 0: 263 | break 264 | 265 | args_help = _parse_args_from_docstring(cls.__init__.__doc__ or cls.__doc__ or "") 266 | 267 | for arg, arg_types, arg_default in args_and_types: 268 | arg_types = tuple(at for at in allowed_types if at in arg_types) 269 | if not arg_types: 270 | # skip argument with not supported type 271 | continue 272 | arg_kwargs: Dict[str, Any] = {} 273 | if bool in arg_types: 274 | arg_kwargs.update(nargs="?", const=True) 275 | # if the only arg type is bool 276 | if len(arg_types) == 1: 277 | use_type: Callable[[str], Union[bool, int, float, str]] = str_to_bool 278 | elif int in arg_types: 279 | use_type = str_to_bool_or_int 280 | elif str in arg_types: 281 | use_type = str_to_bool_or_str 282 | else: 283 | # filter out the bool as we need to use more general 284 | use_type = [at for at in arg_types if at is not bool][0] 285 | else: 286 | use_type = arg_types[0] 287 | 288 | if arg == "gpus" or arg == "tpu_cores": 289 | use_type = _gpus_allowed_type 290 | 291 | # hack for types in (int, float) 292 | if len(arg_types) == 2 and int in set(arg_types) and float in set(arg_types): 293 | use_type = _int_or_float_type 294 | 295 | # hack for track_grad_norm 296 | if arg == "track_grad_norm": 297 | use_type = float 298 | 299 | # hack for precision 300 | if arg == "precision": 301 | use_type = _precision_allowed_type 302 | 303 | parser.add_argument( 304 | f"--{arg}", dest=arg, default=arg_default, type=use_type, help=args_help.get(arg), **arg_kwargs 305 | ) 306 | 307 | if use_argument_group: 308 | return parent_parser 309 | return parser 310 | 311 | 312 | def _parse_args_from_docstring(docstring: str) -> Dict[str, str]: 313 | arg_block_indent = None 314 | current_arg = "" 315 | parsed = {} 316 | for line in docstring.split("\n"): 317 | stripped = line.lstrip() 318 | if not stripped: 319 | continue 320 | line_indent = len(line) - len(stripped) 321 | if stripped.startswith(("Args:", "Arguments:", "Parameters:")): 322 | arg_block_indent = line_indent + 4 323 | elif arg_block_indent is None: 324 | continue 325 | elif line_indent < arg_block_indent: 326 | break 327 | elif line_indent == arg_block_indent: 328 | current_arg, arg_description = stripped.split(":", maxsplit=1) 329 | parsed[current_arg] = arg_description.lstrip() 330 | elif line_indent > arg_block_indent: 331 | parsed[current_arg] += f" {stripped}" 332 | return parsed 333 | 334 | 335 | def _gpus_allowed_type(x: str) -> Union[int, str]: 336 | if "," in x: 337 | return str(x) 338 | return int(x) 339 | 340 | 341 | def _int_or_float_type(x: Union[int, float, str]) -> Union[int, float]: 342 | if "." in str(x): 343 | return float(x) 344 | return int(x) 345 | 346 | 347 | def _precision_allowed_type(x: Union[int, str]) -> Union[int, str]: 348 | """ 349 | >>> _precision_allowed_type("32") 350 | 32 351 | >>> _precision_allowed_type("bf16") 352 | 'bf16' 353 | """ 354 | try: 355 | return int(x) 356 | except ValueError: 357 | return x 358 | 359 | 360 | def _defaults_from_env_vars(fn: Callable) -> Callable: 361 | @wraps(fn) 362 | def insert_env_defaults(self: Any, *args: Any, **kwargs: Any) -> Any: 363 | cls = self.__class__ # get the class 364 | if args: # in case any args passed move them to kwargs 365 | # parse only the argument names 366 | cls_arg_names = [arg[0] for arg in get_init_arguments_and_types(cls)] 367 | # convert args to kwargs 368 | kwargs.update(dict(zip(cls_arg_names, args))) 369 | env_variables = vars(parse_env_variables(cls)) 370 | # update the kwargs by env variables 371 | kwargs = dict(list(env_variables.items()) + list(kwargs.items())) 372 | 373 | # all args were already moved to kwargs 374 | return fn(self, **kwargs) 375 | 376 | return insert_env_defaults -------------------------------------------------------------------------------- /automate/brep.py: -------------------------------------------------------------------------------- 1 | import torch_geometric as tg 2 | import torch 3 | import numpy as np 4 | from dotmap import DotMap 5 | from .conversions import torchify 6 | from torch_geometric.data import Batch 7 | from automate_cpp import Part, PartOptions 8 | import json 9 | import os 10 | import xxhash 11 | import torch_scatter 12 | 13 | class PartDataset(torch.utils.data.Dataset): 14 | def __init__( 15 | self, 16 | splits_path, 17 | data_dir, 18 | mode = 'train', 19 | cache_dir = None, 20 | part_options = None, 21 | graph_options = None 22 | ): 23 | super().__init__() 24 | self.splits_path = splits_path 25 | self.mode = mode 26 | self.cache_dir = cache_dir 27 | self.data_dir = data_dir 28 | 29 | if self.cache_dir is not None: 30 | os.makedirs(os.path.join(self.cache_dir, self.mode), exist_ok=True) 31 | 32 | with open(splits_path, 'r') as f: 33 | splits = json.load(f) 34 | self.part_paths = splits[self.mode] 35 | 36 | self.options = PartOptions() if part_options is None else part_options 37 | 38 | self.features = PartFeatures() if graph_options is None else graph_options 39 | 40 | def __getitem__(self, idx): 41 | if isinstance(idx, slice): 42 | return [self[i] for i in range(len(self.part_paths))[idx]] 43 | if not self.cache_dir is None: 44 | cache_file = os.path.join(self.cache_dir, self.mode, f'{idx}.pt') 45 | if os.path.exists(cache_file): 46 | return torch.load(cache_file) 47 | part_path = os.path.join(self.data_dir, self.part_paths[idx]) 48 | if part_path.endswith('.pt'): 49 | part = torch.load(part_path) 50 | else: 51 | part = Part(part_path, self.options) 52 | graph = part_to_graph(part, self.features) 53 | if not self.cache_dir is None: 54 | torch.save(graph, cache_file) 55 | return graph 56 | 57 | def __len__(self): 58 | return len(self.part_paths) 59 | 60 | # Convert a Part object into a pytorch-geometric compatible 61 | # heterogeneous graph 62 | 63 | # We use a custom torch geometry data object to implement heterogeneous 64 | # graphs since pytorch geometric had not yet implemented them when 65 | # this project was started 66 | 67 | class PartFeatures: 68 | r""" 69 | Options for what to load into a BREP Graph Data Object 70 | """ 71 | def __init__(self): 72 | 73 | # Topology Information 74 | self.brep = True # Include Topology Level Information 75 | 76 | self.face = FaceFeatures() 77 | self.loop = LoopFeatures() 78 | self.edge = EdgeFeatures() 79 | self.vertex = VertexFeatures() 80 | 81 | self.meta_paths = True # Face-Face edge set 82 | 83 | # Include Mesh Data and Relation to BREP topology 84 | self.mesh = True 85 | self.mesh_to_topology = True 86 | 87 | # Include Grid Samples 88 | self.default_num_samples = 10 89 | self.samples = True # Overrides other options 90 | self.face_samples = True 91 | self.normals = True 92 | self.edge_samples = True 93 | self.tangents = True 94 | 95 | # Include random samples 96 | self.random_samples = True 97 | self.num_uniform_samples = 10000 98 | self.uniform_samples = False 99 | 100 | # Part Level Data 101 | self.bounding_box = True 102 | self.volume = True 103 | self.center_of_gravity = True 104 | self.moment_of_inertia = True 105 | self.surface_area = True 106 | 107 | # Mating Coordinate Frames 108 | self.mcfs = True 109 | 110 | def to_flat(x): 111 | r""" 112 | Ensure input is a flat torch float tensor 113 | """ 114 | number = (float, int, bool) 115 | if isinstance(x, number): 116 | return torch.tensor(x).float().reshape((1,)) 117 | if isinstance(x, np.ndarray): 118 | return torch.from_numpy(x).flatten().float() 119 | if torch.is_tensor(x): 120 | return x.float().flatten() 121 | 122 | def to_index(x): 123 | r""" 124 | Ensure input is a torch long tensor (necessary for edge arrays) 125 | """ 126 | if isinstance(x, int): 127 | return torch.tensor(x).long().reshape((1,1)) 128 | if isinstance(x, np.ndarray): 129 | return torch.from_numpy(x).long() 130 | if torch.is_tensor(x): 131 | return x.long() 132 | 133 | def pad_to(tensor, length): 134 | if isinstance(tensor, list): 135 | return pad_to(torch.tensor(tensor), length) 136 | return torch.cat([tensor, torch.zeros(length - tensor.size(0))]).float() 137 | 138 | FACE_PARAM_SIZE = 11 # The maximum parameter size for a face is 8 139 | # so we will pad all parameter arrays to size 8 140 | class FaceFeatures: 141 | r""" 142 | Options for which features to load for each BREP Face 143 | """ 144 | def __init__(self): 145 | 146 | self.parametric_function = True 147 | self.parameter_values = True 148 | self.exclude_origin = False 149 | 150 | self.orientation = True 151 | 152 | self.surface_area = True 153 | self.circumference = True 154 | 155 | 156 | self.bounding_box = True 157 | self.na_bounding_box = True 158 | self.center_of_gravity = True 159 | self.moment_of_inertia = True 160 | 161 | def size(self): 162 | s = 0 163 | if self.parametric_function: 164 | s += 13 165 | if self.parameter_values: 166 | s += FACE_PARAM_SIZE 167 | 168 | if self.exclude_origin: 169 | s -= 3 170 | if self.orientation: 171 | s += 1 172 | 173 | if self.surface_area: 174 | s += 1 175 | if self.circumference: 176 | s += 1 177 | if self.bounding_box: 178 | s += 6 179 | if self.na_bounding_box: 180 | s += 15 181 | if self.center_of_gravity: 182 | s += 3 183 | if self.moment_of_inertia: 184 | s += 9 185 | 186 | return s 187 | 188 | 189 | 190 | def featurize_face(f, options): 191 | if not isinstance(f, dict): 192 | f = DotMap(torchify(f)) 193 | feature_parts = [] 194 | if options.face.parametric_function: 195 | feature_parts.append( 196 | torch.nn.functional.one_hot( 197 | torch.tensor(f.function.value), 198 | f.function.enum_size)) 199 | if options.face.parameter_values: 200 | params = pad_to(f.parameters, FACE_PARAM_SIZE) 201 | if options.face.exclude_origin: 202 | params = params[3:] 203 | feature_parts.append(params) 204 | if options.face.orientation: 205 | feature_parts.append(to_flat(f.orientation)) 206 | 207 | if options.face.surface_area: 208 | feature_parts.append(to_flat(f.surface_area)) 209 | if options.face.circumference: 210 | feature_parts.append(to_flat(f.circumference)) 211 | if options.face.bounding_box: 212 | feature_parts.append(to_flat(f.bounding_box)) 213 | if options.face.na_bounding_box: 214 | feature_parts.append(to_flat(f.na_bounding_box)) 215 | if options.face.center_of_gravity: 216 | feature_parts.append(to_flat(f.center_of_gravity)) 217 | if options.face.moment_of_inertia: 218 | feature_parts.append(to_flat(f.moment_of_inertia)) 219 | 220 | return torch.cat(feature_parts).flatten().float() 221 | 222 | 223 | class LoopFeatures: 224 | r""" 225 | Options for which features to load for each BREP Loop 226 | """ 227 | def __init__(self): 228 | self.type = True 229 | self.length = True 230 | self.na_bounding_box = True 231 | self.center_of_gravity = True 232 | self.moment_of_inertia = True 233 | 234 | def size(self): 235 | s = 0 236 | if self.type: 237 | s += 10 238 | if self.length: 239 | s += 1 240 | if self.na_bounding_box: 241 | s += 15 242 | if self.center_of_gravity: 243 | s += 3 244 | if self.moment_of_inertia: 245 | s += 9 246 | 247 | return s 248 | 249 | 250 | def featurize_loop(l, options): 251 | if not isinstance(l, dict): 252 | l = DotMap(torchify(l)) 253 | feature_parts = [] 254 | if options.loop.type: 255 | feature_parts.append( 256 | torch.nn.functional.one_hot( 257 | torch.tensor(l.type.value), 258 | l.type.enum_size)) 259 | 260 | if options.loop.length: 261 | feature_parts.append(to_flat(l.length)) 262 | if options.loop.na_bounding_box: 263 | feature_parts.append(to_flat(l.na_bounding_box)) 264 | if options.loop.center_of_gravity: 265 | feature_parts.append(to_flat(l.center_of_gravity)) 266 | if options.loop.moment_of_inertia: 267 | feature_parts.append(to_flat(l.moment_of_inertia)) 268 | 269 | return torch.cat(feature_parts).flatten().float() 270 | 271 | 272 | EDGE_PARAM_SIZE = 11 273 | 274 | class EdgeFeatures: 275 | r""" 276 | Options for which features to load for each BREP Edge 277 | """ 278 | def __init__(self): 279 | # Parametric Definition 280 | self.parametric_function = True 281 | self.parameter_values = True 282 | self.exclude_origin = False 283 | self.orientation = True 284 | 285 | self.t_range = True # Parametric Range 286 | 287 | # 3D Start and End and Mid Points 288 | self.start = True 289 | self.end = True 290 | self.mid_point = True 291 | 292 | self.length = True 293 | 294 | self.bounding_box = True 295 | self.na_bounding_box = True 296 | 297 | self.center_of_gravity = True 298 | self.moment_of_inertia = True 299 | 300 | def size(self): 301 | s = 0 302 | if self.parametric_function: 303 | s += 11 304 | if self.parameter_values: 305 | s += EDGE_PARAM_SIZE 306 | 307 | if self.exclude_origin: 308 | s -= 3 309 | if self.orientation: 310 | s += 1 311 | if self.t_range: 312 | s += 2 313 | if self.length: 314 | s += 1 315 | if self.start: 316 | s += 3 317 | if self.end: 318 | s += 3 319 | if self.mid_point: 320 | s += 3 321 | if self.bounding_box: 322 | s += 6 323 | if self.na_bounding_box: 324 | s += 15 325 | if self.center_of_gravity: 326 | s += 3 327 | if self.moment_of_inertia: 328 | s += 9 329 | return s 330 | 331 | 332 | def featurize_edge(e, options): 333 | if not isinstance(e, dict): 334 | e = DotMap(torchify(e)) 335 | feature_parts = [] 336 | if options.edge.parametric_function: 337 | feature_parts.append( 338 | torch.nn.functional.one_hot( 339 | torch.tensor(e.function.value), 340 | e.function.enum_size)) 341 | if options.edge.parameter_values: 342 | params = pad_to(e.parameters, EDGE_PARAM_SIZE) 343 | if options.edge.exclude_origin: 344 | params = params[3:] 345 | feature_parts.append(params) 346 | if options.edge.orientation: 347 | feature_parts.append(to_flat(e.orientation)) 348 | 349 | if options.edge.t_range: 350 | feature_parts.append(to_flat(e.t_range)) 351 | 352 | if options.edge.length: 353 | feature_parts.append(to_flat(e.length)) 354 | 355 | if options.edge.start: 356 | feature_parts.append(to_flat(e.start)) 357 | if options.edge.end: 358 | feature_parts.append(to_flat(e.end)) 359 | if options.edge.mid_point: 360 | feature_parts.append(to_flat(e.mid_point)) 361 | 362 | if options.edge.bounding_box: 363 | feature_parts.append(to_flat(e.bounding_box)) 364 | if options.edge.na_bounding_box: 365 | feature_parts.append(to_flat(e.na_bounding_box)) # TODO - does this actually contain everything? 366 | if options.edge.center_of_gravity: 367 | feature_parts.append(to_flat(e.center_of_gravity)) 368 | if options.edge.moment_of_inertia: 369 | feature_parts.append(to_flat(e.moment_of_inertia)) 370 | 371 | return torch.cat(feature_parts).flatten().float() 372 | 373 | 374 | class VertexFeatures: 375 | r""" 376 | Options for which features to load for each BREP Vertex 377 | """ 378 | def __init__(self): 379 | self.position = True 380 | 381 | def size(self): 382 | return 3 if self.position else 0 383 | 384 | def featurize_vert(v, options): 385 | if not isinstance(v, dict): 386 | v = DotMap(torchify(v)) 387 | feature_parts = [] 388 | if options.vertex.position: 389 | feature_parts.append(to_flat(v.position)) 390 | return torch.cat(feature_parts).flatten().float() 391 | 392 | 393 | def flatbatch(datalist): 394 | follow_batch = [] 395 | if hasattr(datalist[0], 'mcfs'): 396 | follow_batch .append('mcfs') 397 | batch = Batch.from_data_list(datalist, follow_batch=follow_batch) 398 | data = HetData() 399 | for key in dir(batch): 400 | if not key.endswith('batch') and key != 'ptr': 401 | val = getattr(batch, key) 402 | if isinstance(val, torch.Tensor): 403 | setattr(data, key, val) 404 | data.__edge_sets__ = datalist[0].__edge_sets__ 405 | if hasattr(batch, 'mcfs_batch'): 406 | data.mcf_to_graph_idx = batch.mcfs_batch.expand((1, batch.mcfs.shape[0])) 407 | data.__edge_sets__['mcf_to_graph_idx'] = ['graph_idx'] 408 | data.__num_nodes__ = batch.num_nodes 409 | data.__node_sets__ = datalist[0].__node_sets__ 410 | data.__edge_sets__['flat_topos_to_graph_idx'] = ['graph_idx'] 411 | return data 412 | 413 | 414 | def part_to_graph(part, options): 415 | # Add dot (.) access to deserialized parts so they act more like C++ module Parts 416 | if isinstance(part, dict): 417 | part = DotMap(part) 418 | 419 | data = HetData() 420 | 421 | # Keep track of which graph is which during batching 422 | data.graph_idx = to_index(0) 423 | 424 | # It is useful to have a unified "topology" node set for referencing 425 | # arbitrary topological entities against. We essentially "stack" 426 | # The 4 node types in this order [faces, edges, vertices, loops] 427 | # Note that this is different than the normal order, for compatibility 428 | # with older versions that did not consider loops 429 | # Compute the number of each type of node, and the offsets for their 430 | # indices in the global topology list 431 | n_faces = len(part.brep.nodes.faces) 432 | n_edges = len(part.brep.nodes.edges) 433 | n_vertices = len(part.brep.nodes.vertices) 434 | n_loops = len(part.brep.nodes.loops) 435 | data.n_faces = torch.tensor(n_faces).long() 436 | data.n_edges = torch.tensor(n_edges).long() 437 | data.n_vertices = torch.tensor(n_vertices).long() 438 | data.n_loops = torch.tensor(n_loops).long() 439 | 440 | n_topos = n_faces + n_edges + n_vertices + n_loops 441 | face_offset = 0 442 | edge_offset = n_faces 443 | vertex_offset = edge_offset + n_edges 444 | loop_offset = vertex_offset + n_vertices 445 | topo_offsets = [face_offset, edge_offset, vertex_offset, loop_offset] 446 | 447 | # Setup Node Data 448 | if options.brep: 449 | face_features = [featurize_face(f, options) for f in part.brep.nodes.faces] 450 | loop_features = [featurize_loop(f, options) for f in part.brep.nodes.loops] 451 | edge_features = [featurize_edge(f, options) for f in part.brep.nodes.edges] 452 | vert_features = [featurize_vert(f, options) for f in part.brep.nodes.vertices] 453 | 454 | 455 | data.face_export_ids = torch.tensor([xxhash.xxh32(f.export_id).intdigest() for f in part.brep.nodes.faces]).long() 456 | data.loop_export_ids = torch.tensor([xxhash.xxh32(l.export_id).intdigest() for l in part.brep.nodes.loops]).long() 457 | data.edge_export_ids = torch.tensor([xxhash.xxh32(e.export_id).intdigest() for e in part.brep.nodes.edges]).long() 458 | data.vertex_export_ids = torch.tensor([xxhash.xxh32(v.export_id).intdigest() for v in part.brep.nodes.vertices]).long() 459 | 460 | data.__node_sets__.add('face_export_ids') 461 | data.__node_sets__.add('loop_export_ids') 462 | data.__node_sets__.add('edge_export_ids') 463 | data.__node_sets__.add('vertex_export_ids') 464 | 465 | data.faces = torch.stack(face_features) if face_features else torch.empty((0, options.face.size()), dtype=torch.float) 466 | data.__node_sets__.add('faces') 467 | data.loops = torch.stack(loop_features) if loop_features else torch.empty((0, options.loop.size()), dtype=torch.float) 468 | data.edges = torch.stack(edge_features) if edge_features else torch.empty((0, options.edge.size()), dtype=torch.float) 469 | data.vertices = torch.stack(vert_features) if vert_features else torch.empty((0, options.vertex.size()), dtype=torch.float) 470 | 471 | data.face_to_loop = to_index(part.brep.relations.face_to_loop) 472 | data.__edge_sets__['face_to_loop'] = ['faces', 'loops'] 473 | data.loop_to_edge = to_index(part.brep.relations.loop_to_edge) 474 | data.__edge_sets__['loop_to_edge'] = ['loops', 'edges'] 475 | data.edge_to_vertex = to_index(part.brep.relations.edge_to_vertex) 476 | data.__edge_sets__['edge_to_vertex'] = ['edges', 'vertices'] 477 | 478 | if options.meta_paths: 479 | data.face_to_face = to_index(part.brep.relations.face_to_face) 480 | data.__edge_sets__['face_to_face'] = ['faces', 'faces', 'edges'] 481 | 482 | # Add links to the flattened topology list 483 | data.flat_topos = torch.empty((n_topos,0)).float() 484 | data.num_nodes = n_topos 485 | data.flat_topos_to_graph_idx = torch.zeros((1,n_topos)).long() 486 | data.__edge_sets__['flat_topos_to_graph_idx'] = ['graph_idx'] 487 | 488 | data.face_to_flat_topos = torch.stack([ 489 | torch.arange(n_faces).long(), 490 | torch.arange(n_faces).long() + face_offset 491 | ]) 492 | data.__edge_sets__['face_to_flat_topos'] = ['faces', 'flat_topos'] 493 | 494 | data.edge_to_flat_topos = torch.stack([ 495 | torch.arange(n_edges).long(), 496 | torch.arange(n_edges).long() + edge_offset 497 | ]) 498 | data.__edge_sets__['edge_to_flat_topos'] = ['edges', 'flat_topos'] 499 | 500 | data.loop_to_flat_topos = torch.stack([ 501 | torch.arange(n_loops).long(), 502 | torch.arange(n_loops).long() + loop_offset 503 | ]) 504 | data.__edge_sets__['loop_to_flat_topos'] = ['loops', 'flat_topos'] 505 | 506 | data.vertex_to_flat_topos = torch.stack([ 507 | torch.arange(n_vertices).long(), 508 | torch.arange(n_vertices).long() + vertex_offset 509 | ]) 510 | data.__edge_sets__['vertex_to_flat_topos'] = ['vertices', 'flat_topos'] 511 | 512 | if options.mesh: 513 | data.V = torchify(part.mesh.V).float() 514 | data.F = torchify(part.mesh.F).long().T 515 | data.__edge_sets__['F'] = ['V','V','V'] 516 | 517 | num_faces = data.F.size(1) 518 | 519 | # Only make edges into the brep if we have one 520 | if options.mesh_to_topology and options.brep: 521 | data.F_to_faces = torchify(part.mesh_topology.face_to_topology).long().reshape((1,-1)) 522 | data.__edge_sets__['F_to_faces'] = ['faces'] 523 | 524 | # mesh_topology.edge_to_topology is formatted as a #Fx3 matrix 525 | # that can contain -1 wherever there isn't a correspondence (e.g. mesh 526 | # edges in the centers of topological faces). We need to remove these 527 | # and format it as an edge set 528 | # We do this by exanding out into a 3 x (3x #faces) tensor with 529 | # face indices and positions explicit, then filter out the -1s 530 | edge_to_topo = torchify( 531 | part.mesh_topology.edge_to_topology).long().flatten() 532 | face_idx = torch.arange(num_faces).repeat_interleave(3) 533 | idx_in_face = torch.tensor([0,1,2]).long().repeat(num_faces) 534 | E_to_edges = torch.stack([face_idx, idx_in_face, edge_to_topo]) 535 | E_to_edges = E_to_edges[:, (E_to_edges[2] != -1)] 536 | data.E_to_edges = E_to_edges 537 | data.__edge_sets__['E_to_edges'] = ['F', 3, 'faces'] 538 | 539 | # Similar story with vetrices 540 | vert_to_topo = torchify( 541 | part.mesh_topology.point_to_topology).long().flatten() 542 | vert_indices = torch.arange(vert_to_topo.size(0)) 543 | V_to_vertices = torch.stack([vert_indices, vert_to_topo]) 544 | data.V_to_vertices = V_to_vertices[:, (V_to_vertices[1] != -1)] 545 | data.__edge_sets__['V_to_vertices'] = ['V', 'vertices'] 546 | 547 | 548 | 549 | # Setup Part-Level Data 550 | part_feature_list = [] 551 | if options.volume: 552 | part_feature_list.append(to_flat(part.summary.volume)) 553 | if options.surface_area: 554 | part_feature_list.append(to_flat(part.summary.surface_area)) 555 | if options.center_of_gravity: 556 | part_feature_list.append(to_flat(part.summary.center_of_gravity)) 557 | if options.bounding_box: 558 | part_feature_list.append(to_flat(part.summary.bounding_box)) 559 | if options.moment_of_inertia: 560 | part_feature_list.append(to_flat(part.summary.moment_of_inertia.flatten())) 561 | 562 | if len(part_feature_list) > 0: 563 | part_feature = torch.cat(part_feature_list).reshape((1,-1)) 564 | 565 | data.part_feat = part_feature 566 | 567 | # Setup Samples 568 | if options.samples: 569 | if options.face_samples: 570 | samples = part.samples.face_samples 571 | if isinstance(samples, list): 572 | samples = torchify(samples).float() 573 | # Only use normals if the part object has them 574 | has_normals = (samples.size(1) == 9) 575 | if has_normals and not options.normals: 576 | samples = samples[:,[0,1,2,8],:,:] 577 | data.face_samples = samples 578 | data.__node_sets__.add('face_samples') 579 | if options.edge_samples: 580 | samples = part.samples.edge_samples 581 | if isinstance(samples, list): 582 | if samples: 583 | samples = torchify(samples).float() 584 | else: 585 | samples = torch.empty((0, 7, options.default_num_samples)) 586 | # Only use tangents if the part object has them 587 | has_tangents = (samples.size(1) == 7) 588 | if has_tangents and not options.tangents: 589 | samples = samples[:,:3,:] 590 | data.edge_samples = samples 591 | 592 | # Setup Random Samples 593 | if options.random_samples: 594 | random_samples = part.random_samples.samples 595 | if options.uniform_samples: 596 | surface_areas = torch.tensor([face.surface_area for face in part.brep.nodes.faces]) 597 | face_choices = torch.multinomial(surface_areas, options.num_uniform_samples, replacement=True) 598 | face_counts = torch_scatter.scatter_add(torch.ones(face_choices.shape, dtype=torch.long), face_choices, dim_size=len(surface_areas)) 599 | 600 | all_samples = [] 601 | for sample, count in zip(random_samples, face_counts): 602 | sample = torch.from_numpy(sample).float() 603 | sample_filtered = sample[sample[:,6] > 0.5,:6] 604 | if count < sample_filtered.shape[0]: 605 | sample_filtered = sample_filtered[:count] 606 | all_samples.append(sample_filtered) 607 | data.uniform_samples = torch.vstack(all_samples).float() 608 | else: 609 | if isinstance(random_samples, list): 610 | random_samples = torchify(random_samples) 611 | data.random_samples = random_samples 612 | 613 | # Setup MCFs 614 | if options.mcfs: 615 | mcf_origins = [] 616 | mcf_axes = [] 617 | mcf_refs = [] 618 | for mcf in part.default_mcfs: 619 | mcf_origins.append(mcf.origin) 620 | mcf_axes.append(mcf.axis) 621 | 622 | # torchify computes the size of enums, which we want for 623 | # one-hot encoding 624 | axis_ref = mcf.ref.axis_ref 625 | origin_ref = mcf.ref.origin_ref 626 | 627 | axis_tt = axis_ref.reference_type.value 628 | axis_ti = axis_ref.reference_index + topo_offsets[axis_tt] 629 | 630 | origin_tt = origin_ref.reference_type.value 631 | origin_ti = origin_ref.reference_index + topo_offsets[origin_tt] 632 | origin_it = origin_ref.inference_type.value 633 | 634 | mcf_refs.append([axis_ti, origin_ti, origin_it]) 635 | 636 | if len(mcf_axes) > 0 and isinstance(mcf_axes[0], np.ndarray): 637 | mcf_axes = torch.from_numpy(np.stack(mcf_axes)) 638 | mcf_origins = torch.from_numpy(np.stack(mcf_origins)) 639 | else: 640 | mcf_axes = torch.stack(mcf_axes) 641 | mcf_origins = torch.stack(mcf_origins) 642 | data.mcfs = torch.cat([ 643 | mcf_axes, 644 | mcf_origins],1).float() 645 | 646 | data.mcf_refs = torch.tensor(mcf_refs).long().T 647 | data.__edge_sets__['mcf_refs'] = ['flat_topos','flat_topos',0] 648 | 649 | return data 650 | 651 | 652 | class HetData(tg.data.Data): 653 | r""" 654 | An extension of pytorch-geometric's data objects that allows easier 655 | configuration of heterogeneous graphs. Set __edge_sets__ to be a dictionary 656 | where the keys are strings of the edge attribute names (e.g. 'edge_index'), 657 | and the values are the string names of the node data tensors for the 658 | src and dst sides of each edge. Optionally set __node_sets__ to contain 659 | names of node sets (useful for overriding pygeo defaults like 'faces') 660 | """ 661 | 662 | def __init__(self): 663 | super().__init__() 664 | self.__edge_sets__ = {} 665 | self.__node_sets__ = set() 666 | 667 | def __inc__(self, key, value, *args, **kwargs): 668 | if key in self.__edge_sets__: 669 | def get_sizes(nodes): 670 | if isinstance(nodes, int): 671 | return nodes 672 | if isinstance(nodes, str): 673 | if nodes in self.__edge_sets__: 674 | return self[nodes].size(1) 675 | return self[nodes].size(0) 676 | if isinstance(nodes, list) or isinstance(nodes, tuple): 677 | return torch.tensor([[get_sizes(x)] for x in nodes]) 678 | 679 | return get_sizes(self.__edge_sets__[key]) 680 | if key in self.__node_sets__: 681 | return 0 682 | return super().__inc__(key, value) 683 | 684 | def __cat_dim__(self, key, value, *args, **kwargs): 685 | if key in self.__edge_sets__: 686 | return 1 687 | elif key in self.__node_sets__: 688 | return 0 689 | return super().__cat_dim__(key, value) 690 | -------------------------------------------------------------------------------- /automate/conversions.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | # Functions that convert pybind11 objects to serializable formats 5 | 6 | def torchify(obj): 7 | r""" 8 | Convert to nested dictionary of torch tensors for use with 9 | torch.save. Roughly 6.3x file size of original Parasolid .x_t 10 | """ 11 | 12 | if torch.is_tensor(obj): 13 | return obj 14 | 15 | primitive = (int, str, bool, float) 16 | number = (int, float) 17 | if isinstance(obj, primitive): 18 | return obj 19 | if isinstance(obj, np.ndarray): 20 | return torch.from_numpy(obj) 21 | if isinstance(obj, list): 22 | if len(obj) > 0: 23 | if isinstance(obj[0], number): 24 | return torch.tensor(obj) 25 | if isinstance(obj[0], np.ndarray): 26 | return torch.tensor(np.stack(obj)) 27 | if isinstance(obj[0], list): 28 | if len(obj[0]) > 0: 29 | if isinstance(obj[0][0], np.ndarray): 30 | return torch.tensor( 31 | np.stack([ 32 | np.stack(l) for l in obj 33 | ]) 34 | ) 35 | else: 36 | return [[]] 37 | return [torchify(x) for x in obj] 38 | 39 | keys = [x for x in dir(obj) if not x.startswith('__')] 40 | 41 | if 'name' in keys and 'value' in keys: # this is an enum 42 | return { 43 | 'value': getattr(obj, 'value'), 44 | 'name': getattr(obj, 'name'), 45 | 'enum_size': len(keys) - 2 # the other keys are all enum opts 46 | } 47 | 48 | return dict(zip( 49 | keys, 50 | [torchify(getattr(obj, key)) for key in keys] 51 | )) 52 | 53 | 54 | def jsonify(obj): 55 | r""" 56 | Convert to nested dictionary of python primitives for use with 57 | json.dump. Roughly 10x file size of original Parasolid .x_t 58 | """ 59 | primitive = (int, str, bool, float) 60 | if isinstance(obj, primitive): 61 | return obj 62 | if isinstance(obj, np.ndarray): 63 | return obj.tolist() 64 | if isinstance(obj, list): 65 | return [jsonify(x) for x in obj] 66 | 67 | keys = [x for x in dir(obj) if not x.startswith('__')] 68 | 69 | if 'name' in keys and 'value' in keys: # this is an enum 70 | return { 71 | 'value': getattr(obj, 'value'), 72 | 'name': getattr(obj, 'name'), 73 | 'enum_size': len(keys) - 2 # the other keys are all enum opts 74 | } 75 | 76 | return dict(zip( 77 | keys, 78 | [jsonify(getattr(obj, key)) for key in keys] 79 | )) 80 | -------------------------------------------------------------------------------- /automate/eclasses.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from automate_cpp import find_equivalence_classes 3 | import numpy as np 4 | import pandas as pd 5 | 6 | def find_eclasses(items, tolerance): 7 | if len(items) > 100000: 8 | return find_eclasses_multistage(items, tolerance) 9 | if type(items) == np.ndarray: 10 | items = [np.array(item) for item in items.tolist()] 11 | return find_equivalence_classes(items, tolerance) 12 | 13 | def find_eclasses_multistage(items, tolerance): 14 | data = pd.Series([x.tobytes() for x in items]) 15 | if type(items) != np.ndarray: 16 | items = np.array(items) 17 | unique_indices = (~data.duplicated()) 18 | unique_items = items[unique_indices].tolist() 19 | 20 | unique_class = find_equivalence_classes(unique_items, tolerance) 21 | orig_indices = np.arrange(len(items))[unique_indices].to_list() 22 | unique_class_orig = [orig_indices[x] for x in unique_class] 23 | unique_data = data[unique_indices].to_list() 24 | 25 | class_map = dict(zip(unique_data, unique_class_orig)) 26 | 27 | return [class_map[x] for x in data] 28 | 29 | -------------------------------------------------------------------------------- /automate/plot_confusion_matrix.py: -------------------------------------------------------------------------------- 1 | import seaborn as sn 2 | import numpy as np 3 | from matplotlib.figure import Figure 4 | 5 | def plot_confusion_matrix(cm, labels, normalization = (1, 'Actual Avg.'), name='', ax=None): 6 | norm_axis = normalization[0] 7 | cm=cm.cpu().numpy() 8 | if norm_axis >= 0: 9 | totals = cm.sum(norm_axis, keepdims=True) 10 | totals[np.where(totals == 0)] = 1 11 | cm = cm / totals 12 | fmt = '.1%' 13 | if norm_axis < 0: 14 | fmt = 'n' 15 | 16 | fig = Figure(figsize=(8, 8)) 17 | ax = fig.add_subplot() 18 | ax = sn.heatmap(cm, annot=True, fmt=fmt, cmap='PuBu', ax=ax) 19 | name = name + ' ' if len(name) > 0 else name 20 | ax.set_title(f'{name}Confusion Matrix: {normalization[1]}') 21 | ax.set_ylabel('Actual') 22 | ax.set_xlabel('Predicted') 23 | ax.set_xticklabels(labels, rotation=45) 24 | ax.set_yticklabels(labels, rotation=0) 25 | return fig 26 | -------------------------------------------------------------------------------- /automate/pointnet_encoder.py: -------------------------------------------------------------------------------- 1 | from torch.nn import Module 2 | from automate import LinearBlock 3 | import torch 4 | 5 | class PointNetEncoder(Module): 6 | def __init__(self, K=3, layers=(64, 64, 64, 128, 1024)): 7 | super().__init__() 8 | self.encode = LinearBlock(K, *layers) 9 | self.K = K 10 | def forward(self, pc): 11 | pc2 = pc.reshape(-1, self.K) 12 | x = self.encode(pc2) 13 | x = x.reshape(*pc.shape[:-1], -1) 14 | x_p = torch.max(x, dim=-2)[0] 15 | return x, x_p -------------------------------------------------------------------------------- /automate/sbgcn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch_scatter 3 | from torch.nn import Linear, Sequential, ModuleList, BatchNorm1d, Dropout, LeakyReLU, ReLU 4 | import torch_geometric as tg 5 | from .uvnet_encoders import UVNetCurveEncoder, UVNetSurfaceEncoder 6 | 7 | """ 8 | HetData(n_faces=22, n_edges=52, n_vertices=32, n_loops=22, faces=[22, 56], loops=[22, 38], edges=[52, 63], 9 | vertices=[32, 3], face_to_loop=[2, 22], loop_to_edge=[2, 104], edge_to_vertex=[2, 104], face_to_face=[3, 52], 10 | flat_topos=[128, 0], face_to_flat_topos=[2, 22], edge_to_flat_topos=[2, 52], loop_to_flat_topos=[2, 22], 11 | vertex_to_flat_topos=[2, 32], V=[72, 3], F=[3, 140], F_to_faces=[1, 140], E_to_edges=[3, 184], 12 | V_to_vertices=[2, 32], part_feat=[1, 20], face_samples=[22, 9, 10, 10], edge_samples=[52, 7, 10], mcfs=[246, 6], mcf_refs=[3, 246]) 13 | """ 14 | class SBGCN(torch.nn.Module): 15 | def __init__(self, 16 | f_in_width, 17 | l_in_width, 18 | e_in_width, 19 | v_in_width, 20 | out_width, 21 | k, 22 | use_uvnet_features=False, 23 | crv_in_dim=[0, 1, 2, 3, 4, 5], 24 | srf_in_dim=[0, 1, 2, 3, 4, 5, 8], 25 | crv_emb_dim=64, 26 | srf_emb_dim=64, 27 | ): 28 | super().__init__() 29 | self.use_uvnet_features = use_uvnet_features 30 | self.crv_in_dim = crv_in_dim 31 | self.srf_in_dim = srf_in_dim 32 | if use_uvnet_features: 33 | self.crv_emb_dim = crv_emb_dim 34 | self.srf_emb_dim = srf_emb_dim 35 | self.curv_encoder = UVNetCurveEncoder( 36 | in_channels=len(crv_in_dim), output_dims=crv_emb_dim 37 | ) 38 | self.surf_encoder = UVNetSurfaceEncoder( 39 | in_channels=len(srf_in_dim), output_dims=srf_emb_dim 40 | ) 41 | f_in_width += srf_emb_dim 42 | e_in_width += crv_emb_dim 43 | 44 | self.embed_f_in = LinearBlock(f_in_width, out_width) 45 | self.embed_l_in = LinearBlock(l_in_width, out_width) 46 | self.embed_e_in = LinearBlock(e_in_width, out_width) 47 | self.embed_v_in = LinearBlock(v_in_width, out_width) 48 | 49 | self.V2E = BipartiteResMRConv(out_width) 50 | self.E2L = BipartiteResMRConv(out_width) 51 | self.L2F = BipartiteResMRConv(out_width) 52 | 53 | self.ffLayers = ModuleList() 54 | for i in range(k): 55 | self.ffLayers.append(BipartiteResMRConv(out_width)) 56 | 57 | self.F2L = BipartiteResMRConv(out_width) 58 | self.L2E = BipartiteResMRConv(out_width) 59 | self.E2V = BipartiteResMRConv(out_width) 60 | 61 | 62 | def forward(self, data): 63 | x_f = data.faces 64 | x_l = data.loops 65 | x_e = data.edges 66 | x_v = data.vertices 67 | 68 | 69 | # Compute uvnet features 70 | if self.use_uvnet_features: 71 | hidden_srf_feat = self.surf_encoder(data.face_samples[:,self.srf_in_dim,:,:]) 72 | hidden_crv_feat = self.curv_encoder(data.edge_samples[:,self.crv_in_dim,:]) 73 | x_f = torch.cat((x_f, hidden_srf_feat), dim=1) 74 | x_e = torch.cat((x_e, hidden_crv_feat), dim=1) 75 | 76 | 77 | # Apply Input Encoders 78 | x_f = self.embed_f_in(x_f) 79 | x_l = self.embed_l_in(x_l) 80 | x_e = self.embed_e_in(x_e) 81 | x_v = self.embed_v_in(x_v) 82 | 83 | 84 | # Upward Pass ([[1,0]] flips downwards graph edges) 85 | x_e = self.V2E(x_v, x_e, data.edge_to_vertex[[1,0]]) 86 | x_l = self.E2L(x_e, x_l, data.loop_to_edge[[1,0]]) 87 | x_f = self.L2F(x_l, x_f, data.face_to_loop[[1,0]]) 88 | 89 | # Meta-Edge Spine 90 | for conv in self.ffLayers: 91 | x_f = conv(x_f, x_f, data.face_to_face[:2,:]) 92 | 93 | # Downward Pass 94 | x_l = self.F2L(x_f, x_l, data.face_to_loop) 95 | x_e = self.L2E(x_l, x_e, data.loop_to_edge) 96 | x_v = self.E2V(x_e, x_v, data.edge_to_vertex) 97 | 98 | # Flatten Topology Representations 99 | n_topos = x_f.size(0) + x_l.size(0) + x_e.size(0) + x_v.size(0) 100 | n_feats = x_f.size(1) 101 | x_t = torch.zeros((n_topos, n_feats)).type_as(x_f) 102 | 103 | x_t[data.face_to_flat_topos[1]] = x_f[data.face_to_flat_topos[0]] 104 | x_t[data.edge_to_flat_topos[1]] = x_e[data.edge_to_flat_topos[0]] 105 | x_t[data.vertex_to_flat_topos[1]] = x_v[data.vertex_to_flat_topos[0]] 106 | x_t[data.loop_to_flat_topos[1]] = x_l[data.loop_to_flat_topos[0]] 107 | 108 | # Global Pool 109 | x_p = tg.nn.global_max_pool(x_t, data.flat_topos_to_graph_idx.flatten()) 110 | 111 | return x_t, x_p, x_f, x_l, x_e, x_v 112 | 113 | class BipartiteResMRConv(torch.nn.Module): 114 | def __init__(self, width): 115 | super().__init__() 116 | self.mlp = LinearBlock(2*width, width) 117 | 118 | def forward(self, x_src, x_dst, e): 119 | diffs = torch.index_select(x_dst, 0, e[1]) - torch.index_select(x_src, 0, e[0]) 120 | maxes, _ = torch_scatter.scatter_max( 121 | diffs, 122 | e[1], 123 | dim=0, 124 | dim_size=x_dst.shape[0] 125 | ) 126 | return x_dst + self.mlp(torch.cat([x_dst, maxes], dim=1)) 127 | 128 | 129 | class LinearBlock(torch.nn.Module): 130 | def __init__(self, *layer_sizes, batch_norm=False, dropout=0.0, last_linear=False, leaky=True): 131 | super().__init__() 132 | 133 | layers = [] 134 | for i in range(len(layer_sizes) - 1): 135 | c_in = layer_sizes[i] 136 | c_out = layer_sizes[i + 1] 137 | 138 | layers.append(Linear(c_in, c_out)) 139 | if last_linear and i+1 >= len(layer_sizes) - 1: 140 | break 141 | if batch_norm: 142 | layers.append(BatchNorm1d(c_out)) 143 | if dropout > 0: 144 | layers.append(Dropout(p=dropout)) 145 | layers.append((LeakyReLU() if leaky else ReLU())) 146 | 147 | self.f = Sequential(*layers) 148 | 149 | def forward(self, x): 150 | return self.f(x) -------------------------------------------------------------------------------- /automate/util.py: -------------------------------------------------------------------------------- 1 | from argparse import ArgumentParser 2 | import sys 3 | import os 4 | import importlib 5 | import json 6 | from .arg_parsing import from_argparse_args, add_argparse_args 7 | import pytorch_lightning as pl 8 | from pytorch_lightning.loggers import TensorBoardLogger 9 | 10 | class ArgparseInitialized: 11 | @classmethod 12 | def from_argparse_args(cls, args, **kwargs): 13 | return from_argparse_args(cls, args, **kwargs) 14 | 15 | @classmethod 16 | def add_argparse_args(cls, parent_parser): 17 | return add_argparse_args(cls, parent_parser) 18 | 19 | def fix_file_descriptors(): 20 | # When run in a tmux environment, the file_descriptor strategy can run out 21 | # of file handles quickly unless you reset the file handle limit 22 | if sys.platform == "linux" or sys.platform == "linux2": 23 | import resource 24 | from torch.multiprocessing import set_sharing_strategy 25 | set_sharing_strategy("file_descriptor") 26 | resource.setrlimit(resource.RLIMIT_NOFILE, (100_000, 100_000)) 27 | 28 | def get_class(classname: str): 29 | parts = classname.split('.') 30 | module_name = '.'.join(parts[:-1]) 31 | class_name = parts[-1] 32 | module = importlib.import_module(module_name) 33 | return getattr(module, class_name) 34 | 35 | # TODOs: deterministic seeding by default 36 | # 37 | def run_model(default_args = dict()): 38 | fix_file_descriptors() 39 | parser = ArgumentParser(allow_abbrev=False, conflict_handler='resolve') 40 | 41 | parser.add_argument('--argfile', type=str, default=None) 42 | parser.add_argument('--model_class', type=str, default=None) 43 | parser.add_argument('--data_class', type=str, default=None) 44 | 45 | args, _ = parser.parse_known_args() 46 | if args.argfile is not None: 47 | if not os.path.exists(args.argfile): 48 | print(f'Could not find argfile: {args.argfile}') 49 | exit() 50 | with open(args.argfile, 'r') as argfile: 51 | file_args = json.load(argfile) 52 | for key in file_args: 53 | default_args[key] = file_args[key] 54 | default_args['name'] = os.path.splitext(os.path.split(args.argfile)[1])[0] 55 | if 'model_class' in default_args and args.model_class is None: 56 | args.model_class = default_args['model_class'] 57 | if 'data_class' in default_args and args.data_class is None: 58 | args.data_class = default_args['data_class'] 59 | 60 | if args.model_class is None or args.data_class is None: 61 | print('You must specify a --model_class and a --data_class') 62 | exit() 63 | 64 | model_cls = get_class(args.model_class) 65 | data_module_cls = get_class(args.data_class) 66 | 67 | parser = model_cls.add_argparse_args(parser) 68 | parser = data_module_cls.add_argparse_args(parser) 69 | parser = pl.Trainer.add_argparse_args(parser) 70 | 71 | parser.add_argument('--tensorboard_path', type=str, default='.') 72 | parser.add_argument('--checkpoint_path', type=str, default=None) 73 | parser.add_argument('--name', type=str, default='unnamed') 74 | parser.add_argument('--debug', action='store_true') 75 | parser.add_argument('--resume_version', type=int, default=None) 76 | parser.add_argument('--seed', type=int, default=None) 77 | parser.add_argument('--no_train', action='store_true') 78 | parser.add_argument('--no_test', action='store_true') 79 | 80 | parser.set_defaults(**default_args) 81 | 82 | args = parser.parse_args() 83 | 84 | if args.seed is not None: 85 | pl.seed_everything(args.seed) 86 | args.deterministic = True 87 | 88 | #logger = None if args.debug else TensorBoardLogger( 89 | logger = TensorBoardLogger( 90 | args.tensorboard_path, 91 | name=args.name, 92 | default_hp_metric = False, 93 | version=args.resume_version 94 | ) 95 | logger.log_hyperparams(args) 96 | 97 | if not args.debug and args.resume_version is not None and \ 98 | args.checkpoint_path is not None and args.resume_chackpoint is None: 99 | last_ckpt = os.path.join( 100 | os.path.dirname(logger.experiement.log_dir), 101 | 'checkpoints', 102 | 'last.ckpt' 103 | ) 104 | if not os.path.exists(last_ckpt): 105 | print(f'No last checkpoint found for version_{args.resume_version}.') 106 | print(f'Tried {last_ckpt}') 107 | exit() 108 | args.checkpoint_path = last_ckpt 109 | args.resume_from_checkpoint = last_ckpt 110 | 111 | data = data_module_cls.from_argparse_args(args) 112 | if args.checkpoint_path is None: 113 | model = model_cls.from_argparse_args(args) 114 | else: 115 | model = model_cls.load_from_checkpoint(args.checkpoint_path) 116 | 117 | callbacks = model.get_callbacks() 118 | 119 | trainer = pl.Trainer.from_argparse_args(args, logger=logger, callbacks=callbacks) 120 | 121 | if args.auto_lr_find or args.auto_scale_batch_size: 122 | trainer.tune(model, datamodule=data) 123 | 124 | if not args.no_train: 125 | trainer.fit(model, data) 126 | if not args.no_test: 127 | if args.no_train: 128 | if args.checkpoint_path is None: 129 | print('Testing from initialization.') 130 | else: 131 | print(f'Testing from {args.checkpoint_path}') 132 | results = trainer.test(model, datamodule=data) 133 | else: 134 | ckpt = trainer.checkpoint_callback.best_model_path 135 | print(f'Testing from {ckpt}') 136 | results = trainer.test(datamodule=data) -------------------------------------------------------------------------------- /automate/uvnet_encoders.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | import torch.nn.functional as F 4 | from torch.nn.modules.activation import LeakyReLU 5 | 6 | def _conv1d(in_channels, out_channels, kernel_size=3, padding=0, bias=False, batchnorm=True): 7 | """ 8 | Helper function to create a 1D convolutional layer with batchnorm and LeakyReLU activation 9 | 10 | Args: 11 | in_channels (int): Input channels 12 | out_channels (int): Output channels 13 | kernel_size (int, optional): Size of the convolutional kernel. Defaults to 3. 14 | padding (int, optional): Padding size on each side. Defaults to 0. 15 | bias (bool, optional): Whether bias is used. Defaults to False. 16 | 17 | Returns: 18 | nn.Sequential: Sequential contained the Conv1d, BatchNorm1d and LeakyReLU layers 19 | """ 20 | layers = [] 21 | layers.append(nn.Conv1d( in_channels, out_channels, kernel_size=kernel_size, padding=padding, bias=bias )) 22 | if batchnorm: 23 | layers.append(nn.BatchNorm1d(out_channels)) 24 | layers.append(LeakyReLU()) 25 | return nn.Sequential(*layers) 26 | 27 | 28 | def _conv2d(in_channels, out_channels, kernel_size, padding=0, bias=False, batchnorm=True): 29 | """ 30 | Helper function to create a 2D convolutional layer with batchnorm and LeakyReLU activation 31 | 32 | Args: 33 | in_channels (int): Input channels 34 | out_channels (int): Output channels 35 | kernel_size (int, optional): Size of the convolutional kernel. Defaults to 3. 36 | padding (int, optional): Padding size on each side. Defaults to 0. 37 | bias (bool, optional): Whether bias is used. Defaults to False. 38 | 39 | Returns: 40 | nn.Sequential: Sequential contained the Conv2d, BatchNorm2d and LeakyReLU layers 41 | """ 42 | layers = [] 43 | layers.append(nn.Conv2d( 44 | in_channels, 45 | out_channels, 46 | kernel_size=kernel_size, 47 | padding=padding, 48 | bias=bias, )) 49 | if batchnorm: 50 | layers.append(nn.BatchNorm2d(out_channels)) 51 | layers.append(LeakyReLU()) 52 | return nn.Sequential(*layers) 53 | 54 | 55 | def _fc(in_features, out_features, bias=False, batchnorm=True): 56 | layers = [] 57 | layers.append(nn.Linear(in_features, out_features, bias=bias)) 58 | if batchnorm: 59 | layers.append(nn.BatchNorm1d(out_features)) 60 | layers.append(nn.LeakyReLU()) 61 | return nn.Sequential(*layers) 62 | 63 | 64 | class _MLP(nn.Module): 65 | """""" 66 | 67 | def __init__(self, num_layers, input_dim, hidden_dim, output_dim, batchnorm=True): 68 | """ 69 | MLP with linear output 70 | Args: 71 | num_layers (int): The number of linear layers in the MLP 72 | input_dim (int): Input feature dimension 73 | hidden_dim (int): Hidden feature dimensions for all hidden layers 74 | output_dim (int): Output feature dimension 75 | 76 | Raises: 77 | ValueError: If the given number of layers is <1 78 | """ 79 | super(_MLP, self).__init__() 80 | self.linear_or_not = True # default is linear model 81 | self.num_layers = num_layers 82 | self.output_dim = output_dim 83 | self.batchnorm = batchnorm 84 | 85 | if num_layers < 1: 86 | raise ValueError("Number of layers should be positive!") 87 | elif num_layers == 1: 88 | # Linear model 89 | self.linear = nn.Linear(input_dim, output_dim) 90 | else: 91 | # Multi-layer model 92 | self.linear_or_not = False 93 | self.linears = torch.nn.ModuleList() 94 | self.batch_norms = torch.nn.ModuleList() 95 | 96 | self.linears.append(nn.Linear(input_dim, hidden_dim)) 97 | for layer in range(num_layers - 2): 98 | self.linears.append(nn.Linear(hidden_dim, hidden_dim)) 99 | self.linears.append(nn.Linear(hidden_dim, output_dim)) 100 | 101 | # TODO: this could move inside the above loop 102 | for layer in range(num_layers - 1): 103 | self.batch_norms.append(nn.BatchNorm1d((hidden_dim))) 104 | 105 | def forward(self, x): 106 | if self.linear_or_not: 107 | # If linear model 108 | return self.linear(x) 109 | else: 110 | # If MLP 111 | h = x 112 | for i in range(self.num_layers - 1): 113 | h = self.linears[i](h) 114 | if self.batchnorm: 115 | h = self.batch_norms[i](h) 116 | h = F.relu(h) 117 | return self.linears[-1](h) 118 | 119 | 120 | class UVNetCurveEncoder(nn.Module): 121 | def __init__(self, in_channels=6, output_dims=64, batchnorm=False): 122 | """ 123 | This is the 1D convolutional network that extracts features from the B-rep edge 124 | geometry described as 1D UV-grids (see Section 3.2, Curve & surface convolution 125 | in paper) 126 | 127 | Args: 128 | in_channels (int, optional): Number of channels in the edge UV-grids. By default 129 | we expect 3 channels for point coordinates and 3 for 130 | curve tangents. Defaults to 6. 131 | output_dims (int, optional): Output curve embedding dimension. Defaults to 64. 132 | """ 133 | super(UVNetCurveEncoder, self).__init__() 134 | self.in_channels = in_channels 135 | self.conv1 = _conv1d(in_channels, 64, kernel_size=3, padding=1, bias=False, batchnorm=batchnorm) 136 | self.conv2 = _conv1d(64, 128, kernel_size=3, padding=1, bias=False, batchnorm=batchnorm) 137 | self.conv3 = _conv1d(128, 256, kernel_size=3, padding=1, bias=False, batchnorm=batchnorm) 138 | self.final_pool = nn.AdaptiveAvgPool1d(1) 139 | self.fc = _fc(256, output_dims, bias=False, batchnorm=batchnorm) 140 | 141 | for m in self.modules(): 142 | self.weights_init(m) 143 | 144 | def weights_init(self, m): 145 | if isinstance(m, (nn.Linear, nn.Conv1d)): 146 | torch.nn.init.kaiming_uniform_(m.weight.data) 147 | if m.bias is not None: 148 | m.bias.data.fill_(0.0) 149 | 150 | def forward(self, x): 151 | assert x.size(1) == self.in_channels 152 | batch_size = x.size(0) 153 | x = self.conv1(x) 154 | x = self.conv2(x) 155 | x = self.conv3(x) 156 | x = self.final_pool(x) 157 | x = x.view(batch_size, -1) 158 | x = self.fc(x) 159 | return x 160 | 161 | 162 | class UVNetSurfaceEncoder(nn.Module): 163 | def __init__( 164 | self, 165 | in_channels=7, 166 | output_dims=64, 167 | batchnorm=False 168 | ): 169 | """ 170 | This is the 2D convolutional network that extracts features from the B-rep face 171 | geometry described as 2D UV-grids (see Section 3.2, Curve & surface convolution 172 | in paper) 173 | 174 | Args: 175 | in_channels (int, optional): Number of channels in the edge UV-grids. By default 176 | we expect 3 channels for point coordinates and 3 for 177 | surface normals and 1 for the trimming mask. Defaults 178 | to 7. 179 | output_dims (int, optional): Output surface embedding dimension. Defaults to 64. 180 | """ 181 | super(UVNetSurfaceEncoder, self).__init__() 182 | self.in_channels = in_channels 183 | self.conv1 = _conv2d(in_channels, 64, 3, padding=1, bias=False, batchnorm=batchnorm) 184 | self.conv2 = _conv2d(64, 128, 3, padding=1, bias=False, batchnorm=batchnorm) 185 | self.conv3 = _conv2d(128, 256, 3, padding=1, bias=False, batchnorm=batchnorm) 186 | self.final_pool = nn.AdaptiveAvgPool2d(1) 187 | self.fc = _fc(256, output_dims, bias=False, batchnorm=batchnorm) 188 | for m in self.modules(): 189 | self.weights_init(m) 190 | 191 | def weights_init(self, m): 192 | if isinstance(m, (nn.Linear, nn.Conv2d)): 193 | torch.nn.init.kaiming_uniform_(m.weight.data) 194 | if m.bias is not None: 195 | m.bias.data.fill_(0.0) 196 | 197 | def forward(self, x): 198 | assert x.size(1) == self.in_channels 199 | batch_size = x.size(0) 200 | x = self.conv1(x) 201 | x = self.conv2(x) 202 | x = self.conv3(x) 203 | x = self.final_pool(x) 204 | x = x.view(batch_size, -1) 205 | x = self.fc(x) 206 | return x 207 | 208 | 209 | 210 | class _EdgeConv(nn.Module): 211 | def __init__( 212 | self, 213 | edge_feats, 214 | out_feats, 215 | node_feats, 216 | num_mlp_layers=2, 217 | hidden_mlp_dim=64, 218 | ): 219 | """ 220 | This module implements Eq. 2 from the paper where the edge features are 221 | updated using the node features at the endpoints. 222 | 223 | Args: 224 | edge_feats (int): Input edge feature dimension 225 | out_feats (int): Output feature deimension 226 | node_feats (int): Input node feature dimension 227 | num_mlp_layers (int, optional): Number of layers used in the MLP. Defaults to 2. 228 | hidden_mlp_dim (int, optional): Hidden feature dimension in the MLP. Defaults to 64. 229 | """ 230 | super(_EdgeConv, self).__init__() 231 | self.proj = _MLP(1, node_feats, hidden_mlp_dim, edge_feats) 232 | self.mlp = _MLP(num_mlp_layers, edge_feats, hidden_mlp_dim, out_feats) 233 | self.batchnorm = nn.BatchNorm1d(out_feats) 234 | self.eps = torch.nn.Parameter(torch.FloatTensor([0.0])) 235 | 236 | def forward(self, graph, nfeat, efeat): 237 | src, dst = graph.edges() 238 | proj1, proj2 = self.proj(nfeat[src]), self.proj(nfeat[dst]) 239 | agg = proj1 + proj2 240 | h = self.mlp((1 + self.eps) * efeat + agg) 241 | h = F.leaky_relu(self.batchnorm(h)) 242 | return h 243 | 244 | 245 | class _NodeConv(nn.Module): 246 | def __init__( 247 | self, 248 | node_feats, 249 | out_feats, 250 | edge_feats, 251 | num_mlp_layers=2, 252 | hidden_mlp_dim=64, 253 | ): 254 | """ 255 | This module implements Eq. 1 from the paper where the node features are 256 | updated using the neighboring node and edge features. 257 | 258 | Args: 259 | node_feats (int): Input edge feature dimension 260 | out_feats (int): Output feature deimension 261 | node_feats (int): Input node feature dimension 262 | num_mlp_layers (int, optional): Number of layers used in the MLP. Defaults to 2. 263 | hidden_mlp_dim (int, optional): Hidden feature dimension in the MLP. Defaults to 64. 264 | """ 265 | super(_NodeConv, self).__init__() 266 | self.gconv = NNConv( 267 | in_feats=node_feats, 268 | out_feats=out_feats, 269 | edge_func=nn.Linear(edge_feats, node_feats * out_feats), 270 | aggregator_type="sum", 271 | bias=False, 272 | ) 273 | self.batchnorm = nn.BatchNorm1d(out_feats) 274 | self.mlp = _MLP(num_mlp_layers, node_feats, hidden_mlp_dim, out_feats) 275 | self.eps = torch.nn.Parameter(torch.FloatTensor([0.0])) 276 | 277 | def forward(self, graph, nfeat, efeat): 278 | h = (1 + self.eps) * nfeat 279 | h = self.gconv(graph, h, efeat) 280 | h = self.mlp(h) 281 | h = F.leaky_relu(self.batchnorm(h)) 282 | return h 283 | 284 | 285 | class UVNetGraphEncoder(nn.Module): 286 | def __init__( 287 | self, 288 | input_dim, 289 | input_edge_dim, 290 | output_dim, 291 | hidden_dim=64, 292 | learn_eps=True, 293 | num_layers=3, 294 | num_mlp_layers=2, 295 | ): 296 | """ 297 | This is the graph neural network used for message-passing features in the 298 | face-adjacency graph. (see Section 3.2, Message passing in paper) 299 | 300 | Args: 301 | input_dim ([type]): [description] 302 | input_edge_dim ([type]): [description] 303 | output_dim ([type]): [description] 304 | hidden_dim (int, optional): [description]. Defaults to 64. 305 | learn_eps (bool, optional): [description]. Defaults to True. 306 | num_layers (int, optional): [description]. Defaults to 3. 307 | num_mlp_layers (int, optional): [description]. Defaults to 2. 308 | """ 309 | super(UVNetGraphEncoder, self).__init__() 310 | self.num_layers = num_layers 311 | self.learn_eps = learn_eps 312 | 313 | # List of layers for node and edge feature message passing 314 | self.node_conv_layers = torch.nn.ModuleList() 315 | self.edge_conv_layers = torch.nn.ModuleList() 316 | 317 | for layer in range(self.num_layers - 1): 318 | node_feats = input_dim if layer == 0 else hidden_dim 319 | edge_feats = input_edge_dim if layer == 0 else hidden_dim 320 | self.node_conv_layers.append( 321 | _NodeConv( 322 | node_feats=node_feats, 323 | out_feats=hidden_dim, 324 | edge_feats=edge_feats, 325 | num_mlp_layers=num_mlp_layers, 326 | hidden_mlp_dim=hidden_dim, 327 | ), 328 | ) 329 | self.edge_conv_layers.append( 330 | _EdgeConv( 331 | edge_feats=edge_feats, 332 | out_feats=hidden_dim, 333 | node_feats=node_feats, 334 | num_mlp_layers=num_mlp_layers, 335 | hidden_mlp_dim=hidden_dim, 336 | ) 337 | ) 338 | 339 | # Linear function for graph poolings of output of each layer 340 | # which maps the output of different layers into a prediction score 341 | self.linears_prediction = torch.nn.ModuleList() 342 | 343 | for layer in range(num_layers): 344 | if layer == 0: 345 | self.linears_prediction.append(nn.Linear(input_dim, output_dim)) 346 | else: 347 | self.linears_prediction.append(nn.Linear(hidden_dim, output_dim)) 348 | 349 | self.drop1 = nn.Dropout(0.3) 350 | self.drop = nn.Dropout(0.5) 351 | self.pool = MaxPooling() 352 | 353 | def forward(self, g, h, efeat): 354 | hidden_rep = [h] 355 | he = efeat 356 | 357 | for i in range(self.num_layers - 1): 358 | # Update node features 359 | h = self.node_conv_layers[i](g, h, he) 360 | # Update edge features 361 | he = self.edge_conv_layers[i](g, h, he) 362 | hidden_rep.append(h) 363 | 364 | out = hidden_rep[-1] 365 | out = self.drop1(out) 366 | score_over_layer = 0 367 | 368 | # Perform pooling over all nodes in each graph in every layer 369 | for i, h in enumerate(hidden_rep): 370 | pooled_h = self.pool(g, h) 371 | score_over_layer += self.drop(self.linears_prediction[i](pooled_h)) 372 | 373 | return out, score_over_layer 374 | -------------------------------------------------------------------------------- /cpp/automate.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include "eclass.h" 5 | #include "part.h" 6 | 7 | namespace py = pybind11; 8 | using namespace pspy; 9 | 10 | std::string face_repr(const Face& f) { 11 | std::string message = ""; 12 | switch (f.function) { 13 | case SurfaceFunction::PLANE: 14 | message += "PLANE("; 15 | break; 16 | case SurfaceFunction::CYLINDER: 17 | message += "CYLINDER("; 18 | break; 19 | case SurfaceFunction::CONE: 20 | message += "CONE("; 21 | break; 22 | case SurfaceFunction::SPHERE: 23 | message += "SPHERE("; 24 | break; 25 | case SurfaceFunction::TORUS: 26 | message += "TORUS("; 27 | break; 28 | case SurfaceFunction::SPUN: 29 | message += "SPUN("; 30 | break; 31 | case SurfaceFunction::BSURF: 32 | message += "BSURF("; 33 | break; 34 | case SurfaceFunction::OFFSET: 35 | message += "OFFSET("; 36 | break; 37 | case SurfaceFunction::SWEPT: 38 | message += "SWEPT("; 39 | break; 40 | case SurfaceFunction::BLENDSF: 41 | message += "BLENDSF("; 42 | break; 43 | case SurfaceFunction::MESH: 44 | message += "MESH("; 45 | break; 46 | case SurfaceFunction::FSURF: 47 | message += "FSURF"; 48 | break; 49 | case SurfaceFunction::SURFACEOFEXTRUSION: 50 | message += "SURFACEOFEXTRUSION"; 51 | break; 52 | case SurfaceFunction::OTHERSURFACE: 53 | message += "OTHERSURFACE"; 54 | break; 55 | } 56 | for (int i = 0; i < f.parameters.size(); ++i) { 57 | if (i > 0) { 58 | message += ","; 59 | } 60 | message += std::to_string(f.parameters[i]); 61 | } 62 | message += ")"; 63 | return message; 64 | } 65 | 66 | std::string loop_repr(const Loop& l) { 67 | std::string message = ""; 68 | switch (l._type) { 69 | case LoopType::OUTER: 70 | message = ""; 71 | break; 72 | case LoopType::LIKELY_OUTER: 73 | message = ""; 74 | break; 75 | case LoopType::INNER: 76 | message = ""; 77 | break; 78 | case LoopType::LIKELY_INNER: 79 | message = ""; 80 | break; 81 | case LoopType::INNER_SING: 82 | message = ""; 83 | break; 84 | case LoopType::VERTEX: 85 | message = ""; 86 | break; 87 | case LoopType::UNCLEAR: 88 | message = ""; 89 | break; 90 | case LoopType::WINDING: 91 | message = ""; 92 | break; 93 | case LoopType::WIRE: 94 | message = ""; 95 | break; 96 | case LoopType::ERROR: 97 | message = ""; 98 | break; 99 | } 100 | return message; 101 | } 102 | 103 | std::string edge_repr(const Edge& e) { 104 | std::string message = ""; 105 | 106 | switch (e.function) { 107 | case CurveFunction::LINE: 108 | message += "LINE("; 109 | break; 110 | case CurveFunction::CIRCLE: 111 | message += "CIRCLE("; 112 | break; 113 | case CurveFunction::ELLIPSE: 114 | message += "ELLIPSE("; 115 | break; 116 | case CurveFunction::BCURVE: 117 | message += "BCURVE("; 118 | break; 119 | case CurveFunction::ICURVE: 120 | message += "ICURVE("; 121 | break; 122 | case CurveFunction::FCURVE: 123 | message += "FCURVE("; 124 | break; 125 | case CurveFunction::SPCURVE: 126 | message += "SPCURVE("; 127 | break; 128 | case CurveFunction::TRCURVE: 129 | message += "TRCURVE("; 130 | break; 131 | case CurveFunction::CPCURVE: 132 | message += "CPCURVE("; 133 | break; 134 | case CurveFunction::PLINE: 135 | message += "PLINE("; 136 | break; 137 | case CurveFunction::HYPERBOLA: 138 | message += "HYPERBOLA("; 139 | break; 140 | case CurveFunction::PARABOLA: 141 | message += "PARABOLA("; 142 | break; 143 | case CurveFunction::OFFSETCURVE: 144 | message += "OFFSETCURVE("; 145 | break; 146 | case CurveFunction::OTHERCURVE: 147 | message += "OTHERCURVE("; 148 | break; 149 | } 150 | 151 | for (int i = 0; i < e.parameters.size(); ++i) { 152 | if (i > 0) { 153 | message += ","; 154 | } 155 | message += std::to_string(e.parameters[i]); 156 | } 157 | message += ")"; 158 | return message; 159 | } 160 | 161 | PYBIND11_MODULE(automate_cpp, m) { 162 | // part.h 163 | 164 | py::class_(m, "PartOptions") 165 | .def(py::init<>()) 166 | .def_readwrite("just_bb", &PartOptions::just_bb) 167 | .def_readwrite("normalize", &PartOptions::normalize) 168 | .def_readwrite("transform", &PartOptions::transform) 169 | .def_readwrite("transform_matrix", &PartOptions::transform_matrix) 170 | .def_readwrite("num_uv_samples", &PartOptions::num_uv_samples) 171 | .def_readwrite("num_random_samples", &PartOptions::num_random_samples) 172 | .def_readwrite("num_sdf_samples", &PartOptions::num_sdf_samples) 173 | .def_readwrite("sdf_sample_quality", &PartOptions::sdf_sample_quality) 174 | .def_readwrite("sample_normals", &PartOptions::num_uv_samples) 175 | .def_readwrite("sample_tangents", &PartOptions::sample_tangents) 176 | .def_readwrite("tesselate", &PartOptions::tesselate) 177 | .def_readwrite("default_mcfs", &PartOptions::default_mcfs) 178 | .def_readwrite("default_mcfs_only_face_axes", &PartOptions::default_mcfs_only_face_axes) 179 | .def_readwrite("onshape_style", &PartOptions::onshape_style) 180 | .def_readwrite("collect_inferences", &PartOptions::collect_inferences) 181 | .def_readwrite("set_quality", &PartOptions::set_quality) 182 | .def_readwrite("quality", &PartOptions::quality); 183 | 184 | py::class_(m, "Part") 185 | .def(py::init()) 186 | .def(py::init()) 187 | .def_readonly("mesh", &Part::mesh) 188 | .def_readonly("mesh_topology", &Part::mesh_topology) 189 | .def_readonly("brep", &Part::brep) 190 | .def_readonly("samples", &Part::samples) 191 | .def_readonly("random_samples", &Part::random_samples) 192 | .def_readonly("mask_sdf", &Part::mask_sdf) 193 | .def_readonly("summary", &Part::summary) 194 | .def_readonly("inferences", &Part::inferences) 195 | .def_readonly("default_mcfs", &Part::default_mcfs) 196 | .def_readonly("is_valid", &Part::_is_valid); 197 | 198 | py::class_(m, "Mesh") 199 | .def_readonly("V", &Mesh::V) 200 | .def_readonly("F", &Mesh::F); 201 | 202 | py::class_(m, "MeshTopology") 203 | .def_readonly("face_to_topology", &MeshTopology::face_to_topology) 204 | .def_readonly("edge_to_topology", &MeshTopology::edge_to_topology) 205 | .def_readonly("point_to_topology", &MeshTopology::point_to_topology); 206 | 207 | py::class_(m, "InferenceReference") 208 | .def_readonly("reference_index", &InferenceReference::reference_index) 209 | .def_readonly("reference_type", &InferenceReference::reference_type) 210 | .def_readonly("inference_type", &InferenceReference::inference_type); 211 | 212 | py::class_(m, "PartInference") 213 | .def_readonly("origin", &PartInference::origin) 214 | .def_readonly("axis", &PartInference::axis) 215 | .def_readonly("onshape_inference", &PartInference::onshape_inference) 216 | .def_readonly("flipped_in_onshape", &PartInference::flipped_in_onshape) 217 | .def_readonly("reference", &PartInference::reference); 218 | 219 | py::class_(m, "PartFace") 220 | .def_readonly("index", &PartFace::index) 221 | .def_readonly("function", &PartFace::function) 222 | .def_readonly("parameters", &PartFace::parameters) 223 | .def_readonly("orientation", &PartFace::orientation) 224 | .def_readonly("bounding_box", &PartFace::bounding_box) 225 | .def_readonly("na_bounding_box", &PartFace::na_bounding_box) 226 | .def_readonly("surface_area", &PartFace::surface_area) 227 | .def_readonly("circumference", &PartFace::circumference) 228 | .def_readonly("center_of_gravity", &PartFace::center_of_gravity) 229 | .def_readonly("moment_of_inertia", &PartFace::moment_of_inertia) 230 | .def_readonly("loop_neighbors", &PartFace::loop_neighbors) 231 | .def_readonly("edge_neighbors", &PartFace::edge_neighbors) 232 | .def_readonly("vertex_neighbors", &PartFace::vertex_neighbors) 233 | .def_readonly("inferences", &PartFace::inferences) 234 | .def_readonly("export_id", &PartFace::export_id); 235 | 236 | py::class_(m, "PartLoop") 237 | .def_readonly("index", &PartLoop::index) 238 | .def_readonly("type", &PartLoop::type) 239 | .def_readonly("length", &PartLoop::length) 240 | .def_readonly("center_of_gravity", &PartLoop::center_of_gravity) 241 | .def_readonly("moment_of_inertia", &PartLoop::moment_of_inertia) 242 | .def_readonly("na_bounding_box", &PartLoop::na_bounding_box) 243 | .def_readonly("edge_neighbors", &PartLoop::edge_neighbors) 244 | .def_readonly("vertex_neighbors", &PartLoop::vertex_neighbors) 245 | .def_readonly("inferences", &PartLoop::inferences) 246 | .def_readonly("export_id", &PartLoop::export_id); 247 | 248 | py::class_(m, "PartEdge") 249 | .def_readonly("index", &PartEdge::index) 250 | .def_readonly("function", &PartEdge::function) 251 | .def_readonly("parameters", &PartEdge::parameters) 252 | .def_readonly("orientation", &PartEdge::orientation) 253 | .def_readonly("t_range", &PartEdge::t_range) 254 | .def_readonly("start", &PartEdge::start) 255 | .def_readonly("end", &PartEdge::end) 256 | .def_readonly("is_periodic", &PartEdge::is_periodic) 257 | .def_readonly("mid_point", &PartEdge::mid_point) 258 | .def_readonly("length", &PartEdge::length) 259 | .def_readonly("center_of_gravity", &PartEdge::center_of_gravity) 260 | .def_readonly("moment_of_inertia", &PartEdge::moment_of_inertia) 261 | .def_readonly("bounding_box", &PartEdge::bounding_box) 262 | .def_readonly("na_bounding_box", &PartEdge::na_bounding_box) 263 | .def_readonly("vertex_neighbors", &PartEdge::vertex_neighbors) 264 | .def_readonly("inferences", &PartEdge::inferences) 265 | .def_readonly("export_id", &PartEdge::export_id); 266 | 267 | py::class_(m, "PartVertex") 268 | .def_readonly("index", &PartVertex::index) 269 | .def_readonly("position", &PartVertex::position) 270 | .def_readonly("inferences", &PartVertex::inferences) 271 | .def_readonly("export_id", &PartVertex::export_id); 272 | 273 | py::class_(m, "PartTopologyNodes") 274 | .def_readonly("faces", &PartTopologyNodes::faces) 275 | .def_readonly("loops", &PartTopologyNodes::loops) 276 | .def_readonly("edges", &PartTopologyNodes::edges) 277 | .def_readonly("vertices", &PartTopologyNodes::vertices); 278 | 279 | py::class_(m, "PartTopologyRelations") 280 | .def_readonly("face_to_loop", &PartTopologyRelations::face_to_loop) 281 | .def_readonly("loop_to_edge", &PartTopologyRelations::loop_to_edge) 282 | .def_readonly("edge_to_vertex", &PartTopologyRelations::edge_to_vertex) 283 | .def_readonly("face_to_face", &PartTopologyRelations::face_to_face); 284 | 285 | py::class_(m, "PartTopology") 286 | .def_readonly("nodes", &PartTopology::nodes) 287 | .def_readonly("relations", &PartTopology::relations); 288 | 289 | py::class_(m, "PartSamples") 290 | .def_readonly("face_samples", &PartSamples::face_samples) 291 | .def_readonly("edge_samples", &PartSamples::edge_samples); 292 | 293 | py::class_(m, "PartRandomSamples") 294 | .def_readonly("samples", &PartRandomSamples::samples) 295 | .def_readonly("coords", &PartRandomSamples::coords) 296 | .def_readonly("uv_box", &PartRandomSamples::uv_box); 297 | 298 | py::class_(m, "PartMaskSDF") 299 | .def_readonly("coords", &PartMaskSDF::coords) 300 | .def_readonly("sdf", &PartMaskSDF::sdf) 301 | .def_readonly("uv_box", &PartMaskSDF::uv_box); 302 | 303 | py::class_(m, "PartSummary") 304 | .def_readonly("bounding_box", &PartSummary::bounding_box) 305 | .def_readonly("volume", &PartSummary::volume) 306 | .def_readonly("mass", &PartSummary::mass) 307 | .def_readonly("center_of_gravity", &PartSummary::center_of_gravity) 308 | .def_readonly("moment_of_inertia", &PartSummary::moment_of_inertia) 309 | .def_readonly("surface_area", &PartSummary::surface_area) 310 | .def_readonly("topo_type_counts", &PartSummary::topo_type_counts) 311 | .def_readonly("surface_type_counts", &PartSummary::surface_type_counts) 312 | .def_readonly("curve_type_counts", &PartSummary::curve_type_counts) 313 | .def_readonly("loop_type_counts", &PartSummary::loop_type_counts) 314 | .def_readonly("fingerprint", &PartSummary::fingerprint); 315 | 316 | py::class_(m, "PartUniqueInferences") 317 | .def_readonly("origins", &PartUniqueInferences::origins) 318 | .def_readonly("axes", &PartUniqueInferences::axes) 319 | .def_readonly("frames", &PartUniqueInferences::frames) 320 | .def_readonly("origin_references", &PartUniqueInferences::origin_references) 321 | .def_readonly("axes_references", &PartUniqueInferences::axes_references) 322 | .def_readonly("frame_references", &PartUniqueInferences::frame_references); 323 | 324 | py::class_(m, "MCFReference") 325 | .def_readonly("origin_ref", &MCFReference::origin_ref) 326 | .def_readonly("axis_ref", &MCFReference::axis_ref); 327 | 328 | py::class_(m, "MCF") 329 | .def_readonly("origin", &MCF::origin) 330 | .def_readonly("axis", &MCF::axis) 331 | .def_readonly("ref", &MCF::ref); 332 | 333 | // eclass.h 334 | m.def("find_equivalence_classes", &find_equivalence_classes); 335 | 336 | 337 | // body.h 338 | py::class_(m, "Body") 339 | .def("__repr__", 340 | [](const Body& p) { 341 | return ""; 342 | }); 343 | 344 | #ifdef PARASOLID 345 | py::class_(m, "PSBody") 346 | .def("GetTopology", &PSBody::GetTopology) 347 | .def("GetMassProperties", &PSBody::GetMassProperties) 348 | .def("GetBoundingBox", &PSBody::GetBoundingBox) 349 | .def("Transform", &PSBody::Transform) 350 | .def("Tesselate", &PSBody::Tesselate) 351 | .def("__repr__", 352 | [](const PSBody& p) { 353 | return ""; 354 | }); 355 | #endif 356 | 357 | py::class_(m, "OCCTBody") 358 | .def("GetTopology", &OCCTBody::GetTopology) 359 | .def("GetMassProperties", &OCCTBody::GetMassProperties) 360 | .def("GetBoundingBox", &OCCTBody::GetBoundingBox) 361 | .def("Transform", &OCCTBody::Transform) 362 | .def("Tesselate", &OCCTBody::Tesselate) 363 | .def("__repr__", 364 | [](const OCCTBody& p) { 365 | return ""; 366 | }); 367 | 368 | m.def("read_file", &read_file); 369 | #ifdef PARASOLID 370 | m.def("read_xt", &read_xt); 371 | #endif 372 | m.def("read_step", &read_step); 373 | 374 | 375 | 376 | // types.h 377 | 378 | py::enum_(m, "TopologyType") 379 | .value("FACE", TopologyType::FACE) 380 | .value("EDGE", TopologyType::EDGE) 381 | .value("VERTEX", TopologyType::VERTEX) 382 | .value("LOOP", TopologyType::LOOP) 383 | .value("OTHER", TopologyType::OTHER); 384 | 385 | py::enum_(m, "SurfaceFunction") 386 | .value("PLANE", SurfaceFunction::PLANE) 387 | .value("CYLINDER", SurfaceFunction::CYLINDER) 388 | .value("CONE", SurfaceFunction::CONE) 389 | .value("SPHERE", SurfaceFunction::SPHERE) 390 | .value("TORUS", SurfaceFunction::TORUS) 391 | .value("SPUN", SurfaceFunction::SPUN) 392 | .value("BSURF", SurfaceFunction::BSURF) 393 | .value("OFFSET", SurfaceFunction::OFFSET) 394 | .value("SWEPT", SurfaceFunction::SWEPT) 395 | .value("BLENDSF", SurfaceFunction::BLENDSF) 396 | .value("MESH", SurfaceFunction::MESH) 397 | .value("FSURF", SurfaceFunction::FSURF) 398 | .value("SURFACEOFEXTRUSION", SurfaceFunction::SURFACEOFEXTRUSION) 399 | .value("OTHERSURFACE", SurfaceFunction::OTHERSURFACE) 400 | .value("NONE", SurfaceFunction::NONE); 401 | 402 | py::enum_(m, "CurveFunction") 403 | .value("LINE", CurveFunction::LINE) 404 | .value("CIRCLE", CurveFunction::CIRCLE) 405 | .value("ELLIPSE", CurveFunction::ELLIPSE) 406 | .value("BCURVE", CurveFunction::BCURVE) 407 | .value("ICURVE", CurveFunction::ICURVE) 408 | .value("FCURVE", CurveFunction::FCURVE) 409 | .value("SPCURVE", CurveFunction::SPCURVE) 410 | .value("TRCURVE", CurveFunction::TRCURVE) 411 | .value("CPCURVE", CurveFunction::CPCURVE) 412 | .value("PLINE", CurveFunction::PLINE) 413 | .value("HYPERBOLA", CurveFunction::HYPERBOLA) 414 | .value("PARABOLA", CurveFunction::PARABOLA) 415 | .value("OFFSETCURVE", CurveFunction::OFFSETCURVE) 416 | .value("OTHERCURVE", CurveFunction::OTHERCURVE) 417 | .value("NONE", CurveFunction::NONE); 418 | 419 | py::enum_(m, "LoopType") 420 | .value("VERTEX", LoopType::VERTEX) 421 | .value("WIRE", LoopType::WIRE) 422 | .value("OUTER", LoopType::OUTER) 423 | .value("INNER", LoopType::INNER) 424 | .value("WINDING", LoopType::WINDING) 425 | .value("INNER_SING", LoopType::INNER_SING) 426 | .value("LIKELY_OUTER", LoopType::LIKELY_OUTER) 427 | .value("LIKELY_INNER", LoopType::LIKELY_INNER) 428 | .value("UNCLEAR", LoopType::UNCLEAR) 429 | .value("ERROR", LoopType::ERROR); 430 | 431 | py::enum_(m, "InferenceType") 432 | .value("CENTER", InferenceType::CENTER) 433 | .value("CENTROID", InferenceType::CENTROID) 434 | .value("MID_POINT", InferenceType::MID_POINT) 435 | .value("POINT", InferenceType::POINT) 436 | .value("TOP_AXIS_POINT", InferenceType::TOP_AXIS_POINT) 437 | .value("BOTTOM_AXIS_POINT", InferenceType::BOTTOM_AXIS_POINT) 438 | .value("MID_AXIS_POINT", InferenceType::MID_AXIS_POINT) 439 | .value("LOOP_CENTER", InferenceType::LOOP_CENTER); 440 | 441 | py::class_(m, "MassProperties") 442 | .def_readonly("amount", &MassProperties::amount) 443 | .def_readonly("mass", &MassProperties::mass) 444 | .def_readonly("c_of_g", &MassProperties::c_of_g) 445 | .def_readonly("m_of_i", &MassProperties::m_of_i) 446 | .def_readonly("periphery", &MassProperties::periphery) 447 | .def("__repr__", 448 | [](const MassProperties& m) { 449 | return "MassProperties(amount=" + std::to_string(m.amount) + ")"; 450 | }); 451 | 452 | py::class_(m, "Inference") 453 | .def_readonly("z_axis", &Inference::z_axis) 454 | .def_readonly("origin", &Inference::origin) 455 | .def_readonly("inference_type", &Inference::inference_type) 456 | .def_readonly("onshape_inference", &Inference::onshape_inference) 457 | .def_readonly("flipped_in_onshape", &Inference::flipped_in_onshape) 458 | .def("__repr__", 459 | [](const Inference& i) { 460 | return ""; 461 | /* 462 | return "Inference(origin=" + 463 | "(" + std::to_string(i.origin(0)) + "," + std::to_string(i.origin(1)) + "," + std::to_string(i.origin(2)) + ")" + 464 | " , z_axis=" + 465 | "(" + std::to_string(i.origin(0)) + "," + std::to_string(i.origin(1)) + "," + std::to_string(i.origin(2)) + ")" + 466 | ")"; 467 | */ 468 | }); 469 | 470 | 471 | 472 | // topology.h 473 | py::enum_(m, "TopoRelationSense") 474 | .value("None", TopoRelationSense::None) 475 | .value("Negative", TopoRelationSense::Negative) 476 | .value("Positive", TopoRelationSense::Positive); 477 | 478 | py::class_(m, "TopoRelation") 479 | .def_readonly("_parent", &TopoRelation::_parent) 480 | .def_readonly("_child", &TopoRelation::_child) 481 | .def_readonly("_sense", &TopoRelation::_sense) 482 | .def("__repr__", 483 | [](const TopoRelation& tr) { 484 | return "TopoRelation(" + 485 | std::to_string(tr._parent) + "," + 486 | std::to_string(tr._child) + ")"; 487 | }); 488 | 489 | py::class_(m, "BREPTopology") 490 | .def_readonly("faces", &BREPTopology::faces) 491 | .def_readonly("loops", &BREPTopology::loops) 492 | .def_readonly("edges", &BREPTopology::edges) 493 | .def_readonly("vertices", &BREPTopology::vertices) 494 | .def_readonly("face_to_loop", &BREPTopology::face_to_loop) 495 | .def_readonly("loop_to_edge", &BREPTopology::loop_to_edge) 496 | .def_readonly("edge_to_vertex", &BREPTopology::edge_to_vertex) 497 | .def_readonly("loop_to_vertex", &BREPTopology::loop_to_vertex) 498 | .def_readonly("face_to_face", &BREPTopology::face_to_face) 499 | .def_readonly("face_loop", &BREPTopology::face_loop) 500 | .def_readonly("face_edge", &BREPTopology::face_edge) 501 | .def_readonly("face_vertex", &BREPTopology::face_vertex) 502 | .def_readonly("loop_edge", &BREPTopology::loop_edge) 503 | .def_readonly("loop_vertex", &BREPTopology::loop_vertex) 504 | .def_readonly("edge_vertex", &BREPTopology::edge_vertex) 505 | .def_readonly("pk_to_idx", &BREPTopology::pk_to_idx) 506 | .def_readonly("pk_to_class", &BREPTopology::pk_to_class) 507 | .def("__repr__", 508 | [](const BREPTopology& t) { 509 | return ""; 514 | }); 515 | 516 | 517 | // face.h 518 | py::class_(m, "Face") 519 | .def_readonly("function", &Face::function) 520 | .def_readonly("parameters", &Face::parameters) 521 | .def_readonly("orientation", &Face::orientation) 522 | .def_readonly("bounding_box", &Face::bounding_box) 523 | .def_readonly("na_bb_center", &Face::na_bb_center) 524 | .def_readonly("na_bb_x", &Face::na_bb_x) 525 | .def_readonly("na_bb_z", &Face::na_bb_z) 526 | .def_readonly("na_bounding_box", &Face::na_bounding_box) 527 | .def_readonly("surface_area", &Face::surface_area) 528 | .def_readonly("circumference", &Face::circumference) 529 | .def_readonly("center_of_gravity", &Face::center_of_gravity) 530 | .def_readonly("moment_of_inertia", &Face::moment_of_inertia) 531 | .def("__repr__", face_repr); 532 | 533 | #ifdef PARASOLID 534 | py::class_(m, "PSFace") 535 | .def("get_inferences", &PSFace::get_inferences) 536 | .def("sample_points", &PSFace::sample_points) 537 | .def_readonly("function", &PSFace::function) 538 | .def_readonly("parameters", &PSFace::parameters) 539 | .def_readonly("orientation", &PSFace::orientation) 540 | .def_readonly("bounding_box", &PSFace::bounding_box) 541 | .def_readonly("na_bb_center", &PSFace::na_bb_center) 542 | .def_readonly("na_bb_x", &PSFace::na_bb_x) 543 | .def_readonly("na_bb_z", &PSFace::na_bb_z) 544 | .def_readonly("na_bounding_box", &PSFace::na_bounding_box) 545 | .def_readonly("surface_area", &PSFace::surface_area) 546 | .def_readonly("circumference", &PSFace::circumference) 547 | .def_readonly("center_of_gravity", &PSFace::center_of_gravity) 548 | .def_readonly("moment_of_inertia", &PSFace::moment_of_inertia) 549 | .def("__repr__", face_repr); 550 | #endif 551 | 552 | py::class_(m, "OCCTFace") 553 | .def("get_inferences", &OCCTFace::get_inferences) 554 | .def("sample_points", &OCCTFace::sample_points) 555 | .def_readonly("function", &OCCTFace::function) 556 | .def_readonly("parameters", &OCCTFace::parameters) 557 | .def_readonly("orientation", &OCCTFace::orientation) 558 | .def_readonly("bounding_box", &OCCTFace::bounding_box) 559 | .def_readonly("na_bb_center", &OCCTFace::na_bb_center) 560 | .def_readonly("na_bb_x", &OCCTFace::na_bb_x) 561 | .def_readonly("na_bb_z", &OCCTFace::na_bb_z) 562 | .def_readonly("na_bounding_box", &OCCTFace::na_bounding_box) 563 | .def_readonly("surface_area", &OCCTFace::surface_area) 564 | .def_readonly("circumference", &OCCTFace::circumference) 565 | .def_readonly("center_of_gravity", &OCCTFace::center_of_gravity) 566 | .def_readonly("moment_of_inertia", &OCCTFace::moment_of_inertia) 567 | .def("__repr__", face_repr); 568 | 569 | 570 | // loop.h 571 | py::class_(m, "Loop") 572 | .def_readonly("_type", &Loop::_type) 573 | .def_readonly("_is_circle", &Loop::_is_circle) 574 | .def_readonly("length", &Loop::length) 575 | .def_readonly("center_of_gravity", &Loop::center_of_gravity) 576 | .def_readonly("moment_of_inertia", &Loop::moment_of_inertia) 577 | .def_readonly("na_bb_center", &Loop::na_bb_center) 578 | .def_readonly("na_bb_x", &Loop::na_bb_x) 579 | .def_readonly("na_bb_z", &Loop::na_bb_z) 580 | .def_readonly("na_bounding_box", &Loop::na_bounding_box) 581 | .def("__repr__", loop_repr); 582 | 583 | #ifdef PARASOLID 584 | py::class_(m, "PSLoop") 585 | .def("get_inferences", &PSLoop::get_inferences) 586 | .def_readonly("_type", &PSLoop::_type) 587 | .def_readonly("_is_circle", &PSLoop::_is_circle) 588 | .def_readonly("length", &PSLoop::length) 589 | .def_readonly("center_of_gravity", &PSLoop::center_of_gravity) 590 | .def_readonly("moment_of_inertia", &PSLoop::moment_of_inertia) 591 | .def_readonly("na_bb_center", &PSLoop::na_bb_center) 592 | .def_readonly("na_bb_x", &PSLoop::na_bb_x) 593 | .def_readonly("na_bb_z", &PSLoop::na_bb_z) 594 | .def_readonly("na_bounding_box", &PSLoop::na_bounding_box) 595 | .def("__repr__", loop_repr); 596 | #endif 597 | 598 | py::class_(m, "OCCTLoop") 599 | .def("get_inferences", &OCCTLoop::get_inferences) 600 | .def_readonly("_type", &OCCTLoop::_type) 601 | .def_readonly("_is_circle", &OCCTLoop::_is_circle) 602 | .def_readonly("length", &OCCTLoop::length) 603 | .def_readonly("center_of_gravity", &OCCTLoop::center_of_gravity) 604 | .def_readonly("moment_of_inertia", &OCCTLoop::moment_of_inertia) 605 | .def_readonly("na_bb_center", &OCCTLoop::na_bb_center) 606 | .def_readonly("na_bb_x", &OCCTLoop::na_bb_x) 607 | .def_readonly("na_bb_z", &OCCTLoop::na_bb_z) 608 | .def_readonly("na_bounding_box", &OCCTLoop::na_bounding_box) 609 | .def("__repr__", loop_repr); 610 | 611 | // edge.h 612 | py::class_(m, "Edge") 613 | .def_readonly("function", &Edge::function) 614 | .def_readonly("parameters", &Edge::parameters) 615 | .def_readonly("t_start", &Edge::t_start) 616 | .def_readonly("t_end", &Edge::t_end) 617 | .def_readonly("start", &Edge::start) 618 | .def_readonly("end", &Edge::end) 619 | .def_readonly("has_curve", &Edge::_has_curve) 620 | .def_readonly("_is_reversed", &Edge::_is_reversed) 621 | .def_readonly("is_periodic", &Edge::is_periodic) 622 | .def_readonly("mid_point", &Edge::mid_point) 623 | .def_readonly("length", &Edge::length) 624 | .def_readonly("center_of_gravity", &Edge::center_of_gravity) 625 | .def_readonly("moment_of_inertia", &Edge::moment_of_inertia) 626 | .def_readonly("bounding_box", &Edge::bounding_box) 627 | .def_readonly("na_bb_center", &Edge::na_bb_center) 628 | .def_readonly("na_bb_x", &Edge::na_bb_x) 629 | .def_readonly("na_bb_z", &Edge::na_bb_z) 630 | .def_readonly("nn_bounding_box", &Edge::na_bounding_box) 631 | .def("__repr__", edge_repr); 632 | 633 | #ifdef PARASOLID 634 | py::class_(m, "PSEdge") 635 | .def("get_inferences", &PSEdge::get_inferences) 636 | .def("sample_points", &PSEdge::sample_points) 637 | .def_readonly("function", &PSEdge::function) 638 | .def_readonly("parameters", &PSEdge::parameters) 639 | .def_readonly("t_start", &PSEdge::t_start) 640 | .def_readonly("t_end", &PSEdge::t_end) 641 | .def_readonly("start", &PSEdge::start) 642 | .def_readonly("end", &PSEdge::end) 643 | .def_readonly("has_curve", &PSEdge::_has_curve) 644 | .def_readonly("_is_reversed", &PSEdge::_is_reversed) 645 | .def_readonly("is_periodic", &PSEdge::is_periodic) 646 | .def_readonly("mid_point", &PSEdge::mid_point) 647 | .def_readonly("length", &PSEdge::length) 648 | .def_readonly("center_of_gravity", &PSEdge::center_of_gravity) 649 | .def_readonly("moment_of_inertia", &PSEdge::moment_of_inertia) 650 | .def_readonly("bounding_box", &PSEdge::bounding_box) 651 | .def_readonly("na_bb_center", &PSEdge::na_bb_center) 652 | .def_readonly("na_bb_x", &PSEdge::na_bb_x) 653 | .def_readonly("na_bb_z", &PSEdge::na_bb_z) 654 | .def_readonly("nn_bounding_box", &PSEdge::na_bounding_box) 655 | .def("__repr__", edge_repr); 656 | #endif 657 | 658 | py::class_(m, "OCCTEdge") 659 | .def("get_inferences", &OCCTEdge::get_inferences) 660 | .def("sample_points", &OCCTEdge::sample_points) 661 | .def_readonly("function", &OCCTEdge::function) 662 | .def_readonly("parameters", &OCCTEdge::parameters) 663 | .def_readonly("t_start", &OCCTEdge::t_start) 664 | .def_readonly("t_end", &OCCTEdge::t_end) 665 | .def_readonly("start", &OCCTEdge::start) 666 | .def_readonly("end", &OCCTEdge::end) 667 | .def_readonly("has_curve", &OCCTEdge::_has_curve) 668 | .def_readonly("_is_reversed", &OCCTEdge::_is_reversed) 669 | .def_readonly("is_periodic", &OCCTEdge::is_periodic) 670 | .def_readonly("mid_point", &OCCTEdge::mid_point) 671 | .def_readonly("length", &OCCTEdge::length) 672 | .def_readonly("center_of_gravity", &OCCTEdge::center_of_gravity) 673 | .def_readonly("moment_of_inertia", &OCCTEdge::moment_of_inertia) 674 | .def_readonly("bounding_box", &OCCTEdge::bounding_box) 675 | .def_readonly("na_bb_center", &OCCTEdge::na_bb_center) 676 | .def_readonly("na_bb_x", &OCCTEdge::na_bb_x) 677 | .def_readonly("na_bb_z", &OCCTEdge::na_bb_z) 678 | .def_readonly("nn_bounding_box", &OCCTEdge::na_bounding_box) 679 | .def("__repr__", edge_repr); 680 | 681 | 682 | // vertex.h 683 | py::class_(m, "Vertex") 684 | .def_readonly("position", &Vertex::position) 685 | .def("__repr__", 686 | [](const Vertex& v) { 687 | return "Vertex(" + std::to_string(v.position(0)) + "," + std::to_string(v.position(1)) + "," + std::to_string(v.position(2)) + ")"; 688 | }); 689 | 690 | #ifdef PARASOLID 691 | py::class_(m, "PSVertex") 692 | .def("get_inferences", &PSVertex::get_inferences) 693 | .def_readonly("position", &PSVertex::position) 694 | .def("__repr__", 695 | [](const PSVertex& v) { 696 | return "Vertex(" + std::to_string(v.position(0)) + "," + std::to_string(v.position(1)) + "," + std::to_string(v.position(2)) + ")"; 697 | }); 698 | #endif 699 | 700 | 701 | py::class_(m, "OCCTVertex") 702 | .def("get_inferences", &OCCTVertex::get_inferences) 703 | .def_readonly("position", &OCCTVertex::position) 704 | .def("__repr__", 705 | [](const OCCTVertex& v) { 706 | return "Vertex(" + std::to_string(v.position(0)) + "," + std::to_string(v.position(1)) + "," + std::to_string(v.position(2)) + ")"; 707 | }); 708 | 709 | } -------------------------------------------------------------------------------- /cpp/disjointset.cpp: -------------------------------------------------------------------------------- 1 | #include "disjointset.h" 2 | 3 | namespace pspy { 4 | 5 | DisjointSet::DisjointSet(const int size) 6 | { 7 | _id = new int[size]; 8 | for (int i = 0; i < size; ++i) { 9 | _id[i] = -1; 10 | } 11 | } 12 | 13 | DisjointSet::~DisjointSet() 14 | { 15 | delete[] _id; 16 | } 17 | 18 | int DisjointSet::root(int i) 19 | { 20 | while (_id[i] >= 0) { 21 | if (_id[_id[i]] >= 0) { 22 | _id[i] = _id[_id[i]]; 23 | } 24 | i = _id[i]; 25 | } 26 | return i; 27 | } 28 | 29 | void DisjointSet::unite(int p, int q) 30 | { 31 | int i = root(p); 32 | int j = root(q); 33 | if (i == j) return; 34 | if (_id[i] > _id[j]) { 35 | _id[j] += _id[i]; 36 | _id[i] = j; 37 | } 38 | else { 39 | _id[i] += _id[j]; 40 | _id[j] = i; 41 | } 42 | } 43 | 44 | bool DisjointSet::find(int p, int q) 45 | { 46 | return root(p) == root(q); 47 | } 48 | 49 | } -------------------------------------------------------------------------------- /cpp/disjointset.h: -------------------------------------------------------------------------------- 1 | #ifndef DISJOINTSET_INCLUDED 2 | #define DISJOINTSET_INCLUDED 1 3 | 4 | namespace pspy { 5 | 6 | struct DisjointSet 7 | { 8 | public: 9 | DisjointSet(const int size); 10 | ~DisjointSet(); 11 | int root(int i); 12 | void unite(int p, int q); 13 | bool find(int p, int q); 14 | 15 | private: 16 | int* _id; 17 | }; 18 | 19 | } 20 | 21 | #endif // !DISJOINTSET_INCLUDED -------------------------------------------------------------------------------- /cpp/eclass.cpp: -------------------------------------------------------------------------------- 1 | #include "eclass.h" 2 | 3 | #include "lsh.h" 4 | #include "disjointset.h" 5 | #include 6 | 7 | namespace pspy { 8 | 9 | std::vector find_equivalence_classes(const std::vector& points, double tolerance) 10 | { 11 | if (points.size() == 0) { 12 | return std::vector(); 13 | } 14 | int n = points[0].size(); 15 | 16 | auto lsh = LSH(points, n, tolerance); 17 | auto ds = DisjointSet(points.size()); 18 | 19 | for (int i = 0; i < points.size(); ++i) { 20 | assert(points[i].size() == n); 21 | for (int j : lsh.get_nearest_points(points[i])) { 22 | if (i != j) { 23 | ds.unite(i, j); 24 | } 25 | } 26 | } 27 | 28 | std::vector eclasses(points.size()); 29 | for (int i = 0; i < points.size(); ++i) { 30 | eclasses[i] = ds.root(i); 31 | } 32 | 33 | return eclasses; 34 | } 35 | 36 | } -------------------------------------------------------------------------------- /cpp/eclass.h: -------------------------------------------------------------------------------- 1 | #ifndef ECLASS_INCLUDED 2 | #define ECLASS_INCLUDED 1 3 | 4 | #include 5 | #include 6 | 7 | namespace pspy { 8 | 9 | std::vector find_equivalence_classes(const std::vector& points, double tolerance); 10 | 11 | } 12 | 13 | #endif // !ECLASS_INCLUDED 14 | -------------------------------------------------------------------------------- /cpp/lsh.cpp: -------------------------------------------------------------------------------- 1 | #include "lsh.h" 2 | #include 3 | 4 | 5 | namespace pspy { 6 | 7 | long primes[16] = { 8 | 73856093, 9 | 19349669, 10 | 83492791, 11 | 49979539, 12 | 86028157, 13 | 15485867, 14 | 32452843, 15 | 21653669, 16 | 45647201, 17 | 97623107, 18 | 37605409, 19 | 76179877, 20 | 53563061, 21 | 53303309, 22 | 38523587, 23 | 57352609 24 | }; 25 | long computeHashKey(const Eigen::VectorXd& point, int gridIndex, double tolerance) { 26 | assert(point.size() <= 16); // Maximum number of dimensions 27 | long value = 0; 28 | for (int i = 0; i < point.size(); ++i) { 29 | long gridPoint = floor((point(i) + gridIndex * tolerance) / ((point.size() + 1) * tolerance)); 30 | value ^= (primes[i] * gridPoint); 31 | } 32 | return value; 33 | } 34 | 35 | LSH::LSH(const std::vector& points, int nDimensions, double tolerance) { 36 | assert(nDimensions <= 16); // Maximum number of dimensions 37 | tolerance_ = tolerance; 38 | points_ = points; 39 | nDimensions_ = nDimensions; 40 | for (int i = 0; i < (nDimensions_ + 1); ++i) { 41 | IndexHashMap iHashMap; 42 | hashMap_.push_back(iHashMap); 43 | for (size_t index = 0; index < points.size(); ++index) { 44 | assert(points[index].size() == nDimensions_); 45 | long hashValue = computeHashKey(points[index], i, tolerance_); 46 | auto mapElement = hashMap_[i].find(hashValue); 47 | if (mapElement != hashMap_[i].end()) { 48 | (mapElement)->second.insert(index); 49 | } 50 | else { 51 | hashMap_[i][hashValue] = std::set{ (unsigned int)index }; 52 | } 53 | } 54 | } 55 | } 56 | std::set LSH::get_nearest_points(const Eigen::VectorXd& point) const { 57 | std::set nearestPoints{}; 58 | 59 | assert(point.size() == nDimensions_); 60 | for (int i = 0; i < (nDimensions_ + 1); ++i) { 61 | long hashValue = computeHashKey(point, i, tolerance_); 62 | auto mapElement = hashMap_[i].find(hashValue); 63 | if (mapElement != hashMap_[i].end()) 64 | { 65 | for (unsigned int index : mapElement->second) { 66 | if ((points_[index] - point).norm() <= tolerance_) { 67 | nearestPoints.insert(index); 68 | } 69 | } 70 | } 71 | } 72 | return nearestPoints; 73 | } 74 | 75 | } -------------------------------------------------------------------------------- /cpp/lsh.h: -------------------------------------------------------------------------------- 1 | #ifndef LSH_INCLUDED 2 | #define LSH_INCLUDED 1 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | 9 | namespace pspy { 10 | class LSH { 11 | public: 12 | LSH(const std::vector& points, int nDimensions, double tolerance); 13 | // Note, sets can't seem to be passed in as a reference via pybind 14 | std::set get_nearest_points(const Eigen::VectorXd& point) const; 15 | private: 16 | typedef std::unordered_map> IndexHashMap; 17 | std::vector hashMap_; 18 | std::vector points_; 19 | double tolerance_; 20 | int nDimensions_; 21 | }; 22 | 23 | } 24 | 25 | #endif // !LSH_INCLUDED -------------------------------------------------------------------------------- /cpp/part.cpp: -------------------------------------------------------------------------------- 1 | #include "part.h" 2 | #include "eclass.h" 3 | #include 4 | #include 5 | #include 6 | 7 | namespace pspy { 8 | 9 | Part::Part(const std::string& path, PartOptions options) 10 | { 11 | auto bodies = read_file(path); 12 | if (bodies.size() != 1) { 13 | _is_valid = false; 14 | return; 15 | } 16 | _is_valid = true; 17 | auto& body = bodies[0]; 18 | 19 | auto bounding_box = body->GetBoundingBox(); 20 | 21 | if (options.just_bb) { 22 | summary.bounding_box = bounding_box; 23 | return; 24 | } 25 | 26 | if (options.normalize) { 27 | options.transform = true; 28 | Eigen::Vector3d min_corner = bounding_box.row(0); 29 | Eigen::Vector3d max_corner = bounding_box.row(1); 30 | Eigen::Vector3d diag = max_corner - min_corner; 31 | double scale = diag.maxCoeff(); 32 | Eigen::Vector3d center = (max_corner + min_corner) / 2; 33 | options.transform_matrix << 34 | 1, 0, 0, -center(0), 35 | 0, 1, 0, -center(1), 36 | 0, 0, 1, -center(2), 37 | 0, 0, 0, scale; 38 | } 39 | if (options.transform) { 40 | int err = body->Transform(options.transform_matrix); 41 | if (err != 0) { 42 | _is_valid = false; 43 | return; 44 | } 45 | bounding_box = body->GetBoundingBox(); 46 | } 47 | 48 | auto topology = body->GetTopology(); 49 | 50 | auto mass_properties = body->GetMassProperties(); 51 | 52 | if (options.tesselate) { 53 | body->Tesselate( 54 | mesh.V, 55 | mesh.F, 56 | mesh_topology.face_to_topology, 57 | mesh_topology.edge_to_topology, 58 | mesh_topology.point_to_topology, 59 | options.set_quality, 60 | options.quality); 61 | mesh_topology.renumber(topology); 62 | } 63 | 64 | brep.init(topology); 65 | 66 | 67 | if (options.num_uv_samples > 0) { 68 | samples.init(topology, options); 69 | } 70 | 71 | if (options.num_random_samples > 0) { 72 | random_samples.init(topology, options); 73 | } 74 | 75 | if (options.num_sdf_samples > 0) { 76 | mask_sdf.init(topology, options); 77 | } 78 | 79 | summary.init(topology, mass_properties, bounding_box); 80 | 81 | if (options.collect_inferences) { 82 | inferences.init(brep); 83 | } 84 | 85 | if (options.default_mcfs) { 86 | init_default_mcfs(options.onshape_style, options.default_mcfs_only_face_axes); 87 | } 88 | 89 | } 90 | 91 | void Part::init_default_mcfs(bool onshape_style, bool just_face_axes) 92 | { 93 | // Gather Face Oriented MCFs 94 | for (auto& face : brep.nodes.faces) { 95 | for (auto& ax_inf : face.inferences) { 96 | if (ax_inf.onshape_inference || !onshape_style) { 97 | default_mcfs.emplace_back(ax_inf, ax_inf, onshape_style); 98 | // For Axial Types (Cone, Cylinder, Torus, Spun), we don't want 99 | // to duplicate multiple axis references since they are equivalent 100 | // and it does not matter which inference type is used. We will 101 | // arbitrarily choose the top_axis_point since they all have it 102 | if (ax_inf.reference.inference_type == InferenceType::BOTTOM_AXIS_POINT || 103 | ax_inf.reference.inference_type == InferenceType::MID_AXIS_POINT || 104 | ax_inf.reference.inference_type == InferenceType::POINT) { 105 | continue; 106 | } 107 | for (auto& l: face.loop_neighbors) { 108 | auto& loop = brep.nodes.loops[l]; 109 | for (auto& orig_inf : loop.inferences) { 110 | if (orig_inf.onshape_inference || !onshape_style) { 111 | default_mcfs.emplace_back(orig_inf, ax_inf, onshape_style); 112 | } 113 | } 114 | } 115 | 116 | for (auto& e : face.edge_neighbors) { 117 | auto& edge = brep.nodes.edges[e]; 118 | for (auto& orig_inf : edge.inferences) { 119 | if (orig_inf.onshape_inference || !onshape_style) { 120 | default_mcfs.emplace_back(orig_inf, ax_inf, onshape_style); 121 | } 122 | } 123 | } 124 | 125 | for (auto& v : face.vertex_neighbors) { 126 | auto& vert = brep.nodes.vertices[v]; 127 | for (auto& orig_inf : vert.inferences) { 128 | if (orig_inf.onshape_inference || !onshape_style) { 129 | default_mcfs.emplace_back(orig_inf, ax_inf, onshape_style); 130 | } 131 | } 132 | } 133 | 134 | } 135 | } 136 | } 137 | 138 | if (just_face_axes) { 139 | return; 140 | } 141 | 142 | // Gather Loop Oriented MCFs 143 | for (auto& loop : brep.nodes.loops) { 144 | for (auto& ax_inf : loop.inferences) { 145 | if (ax_inf.onshape_inference || !onshape_style) { 146 | default_mcfs.emplace_back(ax_inf, ax_inf, onshape_style); 147 | 148 | for (auto& e : loop.edge_neighbors) { 149 | auto& edge = brep.nodes.edges[e]; 150 | for (auto& orig_inf : edge.inferences) { 151 | if (orig_inf.onshape_inference || !onshape_style) { 152 | default_mcfs.emplace_back(orig_inf, ax_inf, onshape_style); 153 | } 154 | } 155 | } 156 | 157 | for (auto& v : loop.vertex_neighbors) { 158 | auto& vert = brep.nodes.vertices[v]; 159 | for (auto& orig_inf : vert.inferences) { 160 | if (orig_inf.onshape_inference || !onshape_style) { 161 | default_mcfs.emplace_back(orig_inf, ax_inf, onshape_style); 162 | } 163 | } 164 | } 165 | 166 | } 167 | } 168 | } 169 | 170 | // Gather Edge Oriented MCFs 171 | for (auto& edge : brep.nodes.edges) { 172 | for (auto& ax_inf : edge.inferences) { 173 | if (ax_inf.onshape_inference || !onshape_style) { 174 | default_mcfs.emplace_back(ax_inf, ax_inf, onshape_style); 175 | 176 | for (auto& v : edge.vertex_neighbors) { 177 | auto& vert = brep.nodes.vertices[v]; 178 | for (auto& orig_inf : vert.inferences) { 179 | if (orig_inf.onshape_inference || !onshape_style) { 180 | default_mcfs.emplace_back(orig_inf, ax_inf, onshape_style); 181 | } 182 | } 183 | } 184 | 185 | } 186 | } 187 | } 188 | 189 | // Gather Vertex Oriented MCFs 190 | for (auto& vert : brep.nodes.vertices) { 191 | for (auto& ax_inf : vert.inferences) { 192 | if (ax_inf.onshape_inference || !onshape_style) { 193 | default_mcfs.emplace_back(ax_inf, ax_inf, onshape_style); 194 | } 195 | } 196 | } 197 | } 198 | 199 | void MeshTopology::renumber(BREPTopology& topology) 200 | { 201 | for (int i = 0; i < face_to_topology.size(); ++i) { 202 | face_to_topology(i) = topology.pk_to_idx[face_to_topology(i)]; 203 | } 204 | for (int i = 0; i < edge_to_topology.rows(); ++i) { 205 | for (int j = 0; j < edge_to_topology.cols(); ++j) { 206 | 207 | auto topo_pk_idx = edge_to_topology(i, j); 208 | if (topo_pk_idx == 0) { 209 | edge_to_topology(i, j) = -1; 210 | continue; 211 | } 212 | 213 | auto topo_type = topology.pk_to_class[edge_to_topology(i, j)]; 214 | if (topo_type == TopologyType::EDGE) { 215 | edge_to_topology(i, j) = topology.pk_to_idx[topo_pk_idx]; 216 | } 217 | else { 218 | edge_to_topology(i, j) = -1; 219 | } 220 | } 221 | } 222 | for (int i = 0; i < point_to_topology.size(); ++i) { 223 | auto topo_pk_idx = point_to_topology(i); 224 | if (topo_pk_idx == 0) { 225 | point_to_topology(i) = -1; 226 | continue; 227 | } 228 | auto topo_type = topology.pk_to_class[topo_pk_idx]; 229 | if (topo_type == TopologyType::VERTEX) { 230 | point_to_topology(i) = topology.pk_to_idx[topo_pk_idx]; 231 | } 232 | else { 233 | point_to_topology(i) = -1; 234 | } 235 | } 236 | } 237 | 238 | void PartTopology::init(BREPTopology& topology) 239 | { 240 | nodes.init(topology); 241 | relations.init(topology); 242 | } 243 | 244 | void PartTopologyRelations::init(BREPTopology& topology) 245 | { 246 | face_to_loop.resize(2,topology.face_to_loop.size()); 247 | for (int i = 0; i < topology.face_to_loop.size(); ++i) { 248 | face_to_loop(0, i) = topology.face_to_loop[i]._parent; 249 | face_to_loop(1, i) = topology.face_to_loop[i]._child; 250 | } 251 | 252 | loop_to_edge.resize(2, topology.loop_to_edge.size()); 253 | for (int i = 0; i < topology.loop_to_edge.size(); ++i) { 254 | loop_to_edge(0, i) = topology.loop_to_edge[i]._parent; 255 | loop_to_edge(1, i) = topology.loop_to_edge[i]._child; 256 | } 257 | 258 | edge_to_vertex.resize(2, topology.edge_to_vertex.size()); 259 | for (int i = 0; i < topology.edge_to_vertex.size(); ++i) { 260 | edge_to_vertex(0, i) = topology.edge_to_vertex[i]._parent; 261 | edge_to_vertex(1, i) = topology.edge_to_vertex[i]._child; 262 | } 263 | 264 | face_to_face.resize(3, topology.face_to_face.size()); 265 | for (int i = 0; i < topology.face_to_face.size(); ++i) { 266 | face_to_face(0, i) = std::get<0>(topology.face_to_face[i]); 267 | face_to_face(1, i) = std::get<1>(topology.face_to_face[i]); 268 | face_to_face(2, i) = std::get<2>(topology.face_to_face[i]); 269 | } 270 | } 271 | 272 | void PartTopologyNodes::init(BREPTopology& topology) 273 | { 274 | int f = 0; 275 | int l = 0; 276 | int e = 0; 277 | int v = 0; 278 | faces.reserve(topology.faces.size()); 279 | for (auto& face : topology.faces) { 280 | faces.emplace_back(face, f); 281 | ++f; 282 | } 283 | loops.reserve(topology.loops.size()); 284 | for (auto& loop : topology.loops) { 285 | loops.emplace_back(loop, l); 286 | ++l; 287 | } 288 | edges.reserve(topology.edges.size()); 289 | for (auto& edge : topology.edges) { 290 | edges.emplace_back(edge, e); 291 | ++e; 292 | } 293 | vertices.reserve(topology.vertices.size()); 294 | for (auto& vertex : topology.vertices) { 295 | vertices.emplace_back(vertex, v); 296 | ++v; 297 | } 298 | 299 | // Add Downstream Adjacency Lists 300 | for (f = 0; f < faces.size(); ++f) { 301 | faces[f].loop_neighbors = topology.face_loop[f]; 302 | faces[f].edge_neighbors = topology.face_edge[f]; 303 | faces[f].vertex_neighbors = topology.face_vertex[f]; 304 | } 305 | for (l = 0; l < loops.size(); ++l) { 306 | loops[l].edge_neighbors = topology.loop_edge[l]; 307 | loops[l].vertex_neighbors = topology.loop_vertex[l]; 308 | } 309 | for (e = 0; e < edges.size(); ++e) { 310 | edges[e].vertex_neighbors = topology.edge_vertex[e]; 311 | } 312 | 313 | } 314 | 315 | PartFace::PartFace(std::shared_ptr& f, int i) 316 | { 317 | index = i; 318 | function = f->function; 319 | parameters = f->parameters; 320 | orientation = f->orientation; 321 | bounding_box = f->bounding_box; 322 | na_bounding_box.resize(5, 3); 323 | na_bounding_box.block<1, 3>(0, 0) = f->na_bb_center; 324 | na_bounding_box.block<1, 3>(1, 0) = f->na_bb_x; 325 | na_bounding_box.block<1, 3>(2, 0) = f->na_bb_z; 326 | na_bounding_box.block<2, 3>(3, 0) = f->na_bounding_box; 327 | surface_area = f->surface_area; 328 | circumference = f->circumference; 329 | center_of_gravity = f->center_of_gravity; 330 | moment_of_inertia = f->moment_of_inertia; 331 | 332 | auto infs = f->get_inferences(); 333 | inferences.reserve(infs.size()); 334 | for (auto& inf : infs) { 335 | inferences.emplace_back(inf, TopologyType::FACE, index); 336 | } 337 | 338 | export_id = f->export_id; 339 | 340 | } 341 | 342 | PartLoop::PartLoop(std::shared_ptr& l, int i) 343 | { 344 | index = i; 345 | type = l->_type; 346 | length = l->length; 347 | center_of_gravity = l->center_of_gravity; 348 | moment_of_inertia = l->moment_of_inertia; 349 | 350 | na_bounding_box.resize(5, 3); 351 | na_bounding_box.block<1, 3>(0, 0) = l->na_bb_center; 352 | na_bounding_box.block<1, 3>(1, 0) = l->na_bb_x; 353 | na_bounding_box.block<1, 3>(2, 0) = l->na_bb_z; 354 | na_bounding_box.block<2, 3>(3, 0) = l->na_bounding_box; 355 | 356 | auto infs = l->get_inferences(); 357 | inferences.reserve(infs.size()); 358 | for (auto& inf : infs) { 359 | inferences.emplace_back(inf, TopologyType::LOOP, index); 360 | } 361 | 362 | export_id = l->export_id; 363 | } 364 | 365 | PartEdge::PartEdge(std::shared_ptr& e, int i) 366 | { 367 | index = i; 368 | function = e->function; 369 | parameters = e->parameters; 370 | orientation = !e->_is_reversed; 371 | t_range.resize(2); 372 | t_range(0) = e->t_start; 373 | t_range(1) = e->t_end; 374 | start = e->start; 375 | end = e->end; 376 | is_periodic = e->is_periodic; 377 | mid_point = e->mid_point; 378 | length = e->length; 379 | center_of_gravity = e->center_of_gravity; 380 | moment_of_inertia = e->moment_of_inertia; 381 | bounding_box = e->bounding_box; 382 | na_bounding_box.resize(5, 3); 383 | na_bounding_box.block<1, 3>(0, 0) = e->na_bb_center; 384 | na_bounding_box.block<1, 3>(1, 0) = e->na_bb_x; 385 | na_bounding_box.block<1, 3>(2, 0) = e->na_bb_z; 386 | na_bounding_box.block<2, 3>(3, 0) = e->na_bounding_box; 387 | 388 | auto infs = e->get_inferences(); 389 | inferences.reserve(infs.size()); 390 | for (auto& inf : infs) { 391 | inferences.emplace_back(inf, TopologyType::EDGE, index); 392 | } 393 | 394 | export_id = e->export_id; 395 | } 396 | 397 | PartVertex::PartVertex(std::shared_ptr& v, int i) 398 | { 399 | index = i; 400 | position = v->position; 401 | 402 | auto infs = v->get_inferences(); 403 | inferences.reserve(infs.size()); 404 | for (auto& inf : infs) { 405 | inferences.emplace_back(inf, TopologyType::VERTEX, index); 406 | } 407 | 408 | export_id = v->export_id; 409 | } 410 | 411 | void PartSamples::init(BREPTopology& topology, PartOptions options) 412 | { 413 | int num_points = options.num_uv_samples; 414 | bool sample_normals = options.sample_normals; 415 | bool sample_tangents = options.sample_tangents; 416 | Eigen::MatrixXd uv_box; 417 | face_samples.resize(topology.faces.size()); 418 | for (int i = 0; i < topology.faces.size(); ++i) { 419 | topology.faces[i]->sample_points(num_points, sample_normals, face_samples[i], uv_box); 420 | } 421 | 422 | Eigen::Vector2d t_range; 423 | edge_samples.resize(topology.edges.size()); 424 | for (int i = 0; i < topology.edges.size(); ++i) { 425 | topology.edges[i]->sample_points(num_points, sample_tangents, edge_samples[i], t_range); 426 | } 427 | } 428 | 429 | void PartRandomSamples::init(BREPTopology& topology, PartOptions options) 430 | { 431 | const int num_points = options.num_random_samples; 432 | const int n_faces = topology.faces.size(); 433 | samples.resize(n_faces); 434 | coords.resize(n_faces); 435 | uv_box.resize(n_faces); 436 | 437 | 438 | for (int i = 0; i < n_faces; ++i) { 439 | topology.faces[i]->random_sample_points(num_points, samples[i], coords[i], uv_box[i]); 440 | } 441 | } 442 | 443 | void PartMaskSDF::init(BREPTopology& topology, PartOptions options) 444 | { 445 | const int quality = options.sdf_sample_quality; 446 | const int num_points = options.num_sdf_samples; 447 | const int n_faces = topology.faces.size(); 448 | sdf.resize(n_faces); 449 | coords.resize(n_faces); 450 | uv_box.resize(n_faces); 451 | for (int i = 0; i < n_faces; ++i) { 452 | // TODO - commented out so it would compile - implement in the OCCT case, or remove 453 | //topology.faces[i].sample_mask_sdf(quality, num_points, coords[i], sdf[i], uv_box[i]); 454 | } 455 | } 456 | 457 | 458 | 459 | void PartSummary::init(BREPTopology& topology, MassProperties& mass_props, Eigen::MatrixXd& bb) 460 | { 461 | bounding_box = bb; 462 | volume = mass_props.amount; 463 | mass = mass_props.mass; 464 | center_of_gravity = mass_props.c_of_g; 465 | moment_of_inertia = mass_props.m_of_i; 466 | surface_area = mass_props.periphery; 467 | 468 | topo_type_counts.resize(4); 469 | topo_type_counts(0) = topology.faces.size(); 470 | topo_type_counts(1) = topology.edges.size(); 471 | topo_type_counts(2) = topology.vertices.size(); 472 | topo_type_counts(3) = topology.loops.size(); // Put loops last to match old fingerprint 473 | 474 | surface_type_counts.resize(15); 475 | surface_type_counts.setZero(); 476 | for (auto& face : topology.faces) { 477 | int f_idx = static_cast(face->function); 478 | surface_type_counts(f_idx) += 1; 479 | } 480 | 481 | curve_type_counts.resize(15); 482 | curve_type_counts.setZero(); 483 | for (auto& edge : topology.edges) { 484 | int f_idx = static_cast(edge->function); 485 | curve_type_counts(f_idx) += 1; 486 | } 487 | 488 | loop_type_counts.resize(10); 489 | loop_type_counts.setZero(); 490 | for (auto& loop : topology.loops) { 491 | int t_idx = static_cast(loop->_type); 492 | loop_type_counts(t_idx) += 1; 493 | } 494 | 495 | // Collect a hash data vector for part deduplication 496 | fingerprint.resize(14); 497 | fingerprint(0) = volume; 498 | // Take upper triangular half of moment_of_inertia matrix 499 | fingerprint.block<3, 1>(1, 0) = moment_of_inertia.block<1, 3>(0, 0); 500 | fingerprint.block<2, 1>(4, 0) = moment_of_inertia.block<1, 2>(1, 1); 501 | fingerprint(6) = moment_of_inertia(2, 2); 502 | fingerprint.block<3, 1>(7, 0) = center_of_gravity; 503 | fingerprint.block<4, 1>(10,0) = topo_type_counts.block<4,1>(0,0); 504 | } 505 | 506 | PartInference::PartInference(const Inference& inf, TopologyType ref_type, int ref_index) 507 | { 508 | origin = inf.origin; 509 | axis = inf.z_axis; 510 | onshape_inference = inf.onshape_inference; 511 | flipped_in_onshape = inf.flipped_in_onshape; 512 | reference.inference_type = inf.inference_type; 513 | reference.reference_type = ref_type; 514 | reference.reference_index = ref_index; 515 | } 516 | 517 | void PartUniqueInferences::init(const PartTopology& topo, double tolerance) 518 | { 519 | std::vector all_origins; 520 | std::vector all_axes; 521 | std::vector all_frames; 522 | std::vector all_refs; 523 | 524 | // Collect All Inferences for Deduplication 525 | for (auto& t : topo.nodes.faces) { 526 | for (auto& inf : t.inferences) { 527 | all_origins.push_back(inf.origin); 528 | all_axes.push_back(inf.axis); 529 | Eigen::VectorXd frame(6); 530 | frame.block<3, 1>(0, 0) = inf.origin; 531 | frame.block<3, 1>(3, 0) = inf.axis; 532 | all_frames.push_back(frame); 533 | all_refs.push_back(inf.reference); 534 | } 535 | } 536 | for (auto& t : topo.nodes.loops) { 537 | for (auto& inf : t.inferences) { 538 | all_origins.push_back(inf.origin); 539 | all_axes.push_back(inf.axis); 540 | Eigen::VectorXd frame(6); 541 | frame.block<3, 1>(0, 0) = inf.origin; 542 | frame.block<3, 1>(3, 0) = inf.axis; 543 | all_frames.push_back(frame); 544 | all_refs.push_back(inf.reference); 545 | } 546 | } 547 | for (auto& t : topo.nodes.edges) { 548 | for (auto& inf : t.inferences) { 549 | all_origins.push_back(inf.origin); 550 | all_axes.push_back(inf.axis); 551 | Eigen::VectorXd frame(6); 552 | frame.block<3, 1>(0, 0) = inf.origin; 553 | frame.block<3, 1>(3, 0) = inf.axis; 554 | all_frames.push_back(frame); 555 | all_refs.push_back(inf.reference); 556 | } 557 | } 558 | for (auto& t : topo.nodes.vertices) { 559 | for (auto& inf : t.inferences) { 560 | all_origins.push_back(inf.origin); 561 | all_axes.push_back(inf.axis); 562 | Eigen::VectorXd frame(6); 563 | frame.block<3, 1>(0, 0) = inf.origin; 564 | frame.block<3, 1>(3, 0) = inf.axis; 565 | all_frames.push_back(frame); 566 | all_refs.push_back(inf.reference); 567 | } 568 | } 569 | 570 | // Deduplicate inferences 571 | auto origin_eclasses = find_equivalence_classes(all_origins, tolerance); 572 | auto axes_eclasses = find_equivalence_classes(all_axes, tolerance); 573 | auto frame_eclasses = find_equivalence_classes(all_frames, tolerance); 574 | 575 | // Find unique equivalence class ids 576 | std::set origin_eclass_ids(origin_eclasses.begin(), origin_eclasses.end()); 577 | std::set axes_eclass_ids(axes_eclasses.begin(), axes_eclasses.end()); 578 | std::set frame_eclass_ids(frame_eclasses.begin(), frame_eclasses.end()); 579 | 580 | int n_origins = origin_eclass_ids.size(); 581 | int n_axes = axes_eclass_ids.size(); 582 | int n_frames = frame_eclass_ids.size(); 583 | 584 | origins.resize(n_origins, 3); 585 | origin_references.resize(n_origins); 586 | axes.resize(n_axes, 3); 587 | axes_references.resize(n_axes); 588 | frames.resize(n_frames, 6); 589 | frame_references.resize(n_frames); 590 | 591 | // Renumber starting at 0 592 | // Also, copy the unique vectors 593 | int i = 0; 594 | std::map origin_reindex; 595 | for (auto c : origin_eclass_ids) { 596 | origin_reindex[c] = i; 597 | origins.block<1, 3>(i, 0) = all_origins[c]; 598 | ++i; 599 | } 600 | 601 | 602 | i = 0; 603 | std::map axes_reindex; 604 | for (auto c : axes_eclass_ids) { 605 | axes_reindex[c] = i; 606 | axes.block<1, 3>(i, 0) = all_axes[c]; 607 | ++i; 608 | } 609 | 610 | i = 0; 611 | std::map frame_reindex; 612 | for (auto c : frame_eclass_ids) { 613 | frame_reindex[c] = i; 614 | frames.block<1, 6>(i, 0) = all_frames[c]; 615 | ++i; 616 | } 617 | 618 | // Finally, collect references 619 | for (i = 0; i < all_refs.size(); ++i) { 620 | origin_references[origin_reindex[origin_eclasses[i]]].push_back(all_refs[i]); 621 | axes_references[axes_reindex[axes_eclasses[i]]].push_back(all_refs[i]); 622 | frame_references[frame_reindex[frame_eclasses[i]]].push_back(all_refs[i]); 623 | } 624 | } 625 | 626 | MCF::MCF(const PartInference& origin_inf, const PartInference& axis_inf, bool onshape_style) 627 | { 628 | origin = origin_inf.origin; 629 | axis = axis_inf.axis; 630 | if (onshape_style && axis_inf.flipped_in_onshape) { 631 | axis = -axis; 632 | } 633 | ref.origin_ref = origin_inf.reference; 634 | ref.axis_ref = axis_inf.reference; 635 | } 636 | 637 | } -------------------------------------------------------------------------------- /cpp/part.h: -------------------------------------------------------------------------------- 1 | #ifndef PART_H_INCLUDED 2 | #define PART_H_INCLUDED 3 | 4 | #include 5 | #include 6 | #include 7 | #include "body.h" 8 | 9 | namespace pspy { 10 | struct PartOptions { 11 | bool just_bb = false; 12 | bool normalize = false; 13 | bool transform = false; 14 | Eigen::Matrix transform_matrix; 15 | int num_uv_samples = 10; 16 | int num_random_samples = 0; 17 | int num_sdf_samples = 0; 18 | int sdf_sample_quality = 5000; 19 | bool sample_normals = true; 20 | bool sample_tangents = true; 21 | bool tesselate = true; 22 | bool default_mcfs = true; 23 | bool default_mcfs_only_face_axes = false; 24 | bool onshape_style = true; 25 | bool collect_inferences = false; 26 | bool set_quality = false; 27 | double quality = 0.01; 28 | }; 29 | 30 | struct Mesh { 31 | Eigen::MatrixXd V; 32 | Eigen::MatrixXi F; 33 | }; 34 | 35 | struct MeshTopology { 36 | Eigen::VectorXi face_to_topology; 37 | Eigen::MatrixXi edge_to_topology; 38 | Eigen::VectorXi point_to_topology; 39 | 40 | void renumber(BREPTopology& topology); 41 | }; 42 | 43 | struct InferenceReference { 44 | int reference_index; 45 | TopologyType reference_type; 46 | InferenceType inference_type; 47 | }; 48 | 49 | struct PartInference { 50 | PartInference(const Inference& inf, TopologyType ref_type, int ref_index); 51 | Eigen::Vector3d origin; 52 | Eigen::Vector3d axis; 53 | 54 | bool onshape_inference; 55 | bool flipped_in_onshape; 56 | 57 | InferenceReference reference; 58 | }; 59 | 60 | struct PartFace { 61 | PartFace(std::shared_ptr& f, int i); 62 | int index; 63 | SurfaceFunction function; 64 | std::vector parameters; 65 | bool orientation; // True is face normal matches surface normal 66 | Eigen::MatrixXd bounding_box; 67 | Eigen::MatrixXd na_bounding_box; 68 | double surface_area; 69 | double circumference; 70 | Eigen::Vector3d center_of_gravity; 71 | Eigen::MatrixXd moment_of_inertia; 72 | 73 | std::vector loop_neighbors; 74 | std::vector edge_neighbors; 75 | std::vector vertex_neighbors; 76 | 77 | std::vector inferences; 78 | 79 | std::string export_id; 80 | }; 81 | 82 | struct PartLoop { 83 | PartLoop(std::shared_ptr& l, int i); 84 | int index; 85 | LoopType type; 86 | double length; 87 | Eigen::Vector3d center_of_gravity; 88 | Eigen::MatrixXd moment_of_inertia; 89 | Eigen::MatrixXd na_bounding_box; 90 | 91 | std::vector edge_neighbors; 92 | std::vector vertex_neighbors; 93 | 94 | std::vector inferences; 95 | 96 | std::string export_id; 97 | }; 98 | 99 | struct PartEdge { 100 | PartEdge(std::shared_ptr& e, int i); 101 | int index; 102 | CurveFunction function; 103 | std::vector parameters; 104 | bool orientation; // true is edge direction matches curve direction 105 | Eigen::VectorXd t_range; 106 | Eigen::Vector3d start; 107 | Eigen::Vector3d end; 108 | bool is_periodic; 109 | Eigen::Vector3d mid_point; 110 | double length; 111 | Eigen::Vector3d center_of_gravity; 112 | Eigen::MatrixXd moment_of_inertia; 113 | Eigen::MatrixXd bounding_box; 114 | Eigen::MatrixXd na_bounding_box; 115 | 116 | std::vector vertex_neighbors; 117 | 118 | std::vector inferences; 119 | 120 | std::string export_id; 121 | }; 122 | 123 | struct PartVertex { 124 | PartVertex(std::shared_ptr& v, int i); 125 | int index; 126 | Eigen::Vector3d position; 127 | 128 | std::vector inferences; 129 | 130 | std::string export_id; 131 | }; 132 | 133 | struct PartTopologyNodes { 134 | void init(BREPTopology& topology); 135 | std::vector faces; 136 | std::vector loops; 137 | std::vector edges; 138 | std::vector vertices; 139 | }; 140 | 141 | struct PartTopologyRelations { 142 | void init(BREPTopology& topology); 143 | Eigen::MatrixXi face_to_loop; 144 | Eigen::MatrixXi loop_to_edge; 145 | Eigen::MatrixXi edge_to_vertex; 146 | Eigen::MatrixXi face_to_face; // mx3 (third is edge ref) 147 | }; 148 | 149 | struct PartTopology { 150 | void init(BREPTopology& topology); 151 | PartTopologyNodes nodes; 152 | PartTopologyRelations relations; 153 | }; 154 | 155 | struct PartSamples { 156 | void init(BREPTopology& topology, PartOptions options); 157 | std::vector< std::vector< Eigen::MatrixXd> > face_samples; 158 | std::vector< std::vector< Eigen::VectorXd> > edge_samples; 159 | }; 160 | 161 | struct PartRandomSamples { 162 | void init(BREPTopology& topology, PartOptions options); 163 | std::vector samples; 164 | std::vector coords; 165 | std::vector uv_box; 166 | }; 167 | 168 | struct PartMaskSDF { 169 | void init(BREPTopology& topology, PartOptions options); 170 | std::vector coords; 171 | std::vector sdf; 172 | std::vector uv_box; 173 | }; 174 | 175 | struct PartSummary { 176 | void init(BREPTopology& topology, MassProperties& mass_props, Eigen::MatrixXd& bounding_box); 177 | Eigen::MatrixXd bounding_box; 178 | double volume; 179 | double mass; 180 | Eigen::Vector3d center_of_gravity; 181 | Eigen::MatrixXd moment_of_inertia; 182 | double surface_area; 183 | Eigen::VectorXd topo_type_counts; // faces, edges, vertices, loops 184 | Eigen::VectorXd surface_type_counts; 185 | Eigen::VectorXd curve_type_counts; 186 | Eigen::VectorXd loop_type_counts; 187 | Eigen::VectorXd fingerprint; 188 | 189 | }; 190 | 191 | struct PartUniqueInferences { 192 | void init(const PartTopology& topo, double tolerance = 1e-8); 193 | Eigen::MatrixXd origins; 194 | Eigen::MatrixXd axes; 195 | Eigen::MatrixXd frames; 196 | std::vector > origin_references; 197 | std::vector > axes_references; 198 | std::vector > frame_references; 199 | }; 200 | 201 | struct MCFReference { 202 | InferenceReference origin_ref; 203 | InferenceReference axis_ref; 204 | }; 205 | 206 | struct MCF { 207 | MCF(const PartInference& origin_inf, 208 | const PartInference& axis_inf, 209 | bool onshape_style = true // flip axes where onshape would 210 | ); 211 | Eigen::Vector3d origin; 212 | Eigen::Vector3d axis; 213 | MCFReference ref; 214 | }; 215 | 216 | struct Part { 217 | Part(const std::string& path, PartOptions options = PartOptions()); 218 | void init_default_mcfs( 219 | bool onshape_style, // only use onshape specific MCFs 220 | bool just_face_axes // limit to MCFs with face-based axes 221 | ); 222 | Mesh mesh; 223 | MeshTopology mesh_topology; 224 | PartTopology brep; 225 | PartSamples samples; 226 | PartRandomSamples random_samples; 227 | PartMaskSDF mask_sdf; 228 | PartSummary summary; 229 | PartUniqueInferences inferences; 230 | std::vector default_mcfs; 231 | bool _is_valid; 232 | }; 233 | 234 | } 235 | 236 | #endif // !PART_H_INCLUDED 237 | -------------------------------------------------------------------------------- /environment.yml: -------------------------------------------------------------------------------- 1 | name: automate 2 | channels: 3 | - pytorch 4 | - pyg 5 | - fvcore 6 | - iopath 7 | - bottler 8 | - pytorch3d 9 | - conda-forge 10 | - defaults 11 | 12 | dependencies: 13 | - python 14 | - cudatoolkit=11.3 15 | - pytorch>=1.10 16 | - pyg=2.0.3 17 | - numpy 18 | - pandas 19 | - pytorch-lightning 20 | - tensorboard 21 | - pyarrow 22 | - requests 23 | - scipy 24 | - matplotlib 25 | - dotmap 26 | - h5py 27 | - meshplot 28 | - ipykernel 29 | - pip 30 | - eigen 31 | - occt>=7.6 32 | - cmake 33 | - pybind11 34 | - pip: 35 | - pyrender 36 | - xxhash 37 | -------------------------------------------------------------------------------- /minimal_env.yml: -------------------------------------------------------------------------------- 1 | name: automate 2 | channels: 3 | - pytorch 4 | - pyg 5 | - conda-forge 6 | - defaults 7 | 8 | dependencies: 9 | - python 10 | - cudatoolkit=11.3 11 | - pytorch>=1.10 12 | - pyg=2.0.3 13 | - numpy 14 | - pytorch-lightning 15 | - matplotlib 16 | - dotmap 17 | - eigen 18 | - occt>=7.6 19 | - cmake 20 | - pybind11 21 | - seaborn 22 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages, Extension 2 | from setuptools.command.build_ext import build_ext as build_ext_orig 3 | import os 4 | import pathlib 5 | 6 | ## From https://stackoverflow.com/questions/42585210/extending-setuptools-extension-to-use-cmake-in-setup-py ## 7 | class CMakeExtension(Extension): 8 | 9 | def __init__(self, name): 10 | # don't invoke the original build_ext for this special extension 11 | super().__init__(name, sources=[]) 12 | 13 | 14 | class build_ext(build_ext_orig): 15 | 16 | def run(self): 17 | for ext in self.extensions: 18 | self.build_cmake(ext) 19 | super().run() 20 | 21 | def build_cmake(self, ext): 22 | cwd = pathlib.Path().absolute() 23 | 24 | # these dirs will be created in build_py, so if you don't have 25 | # any python sources to bundle, the dirs will be missing 26 | build_temp = pathlib.Path(self.build_temp) 27 | build_temp.mkdir(parents=True, exist_ok=True) 28 | extdir = pathlib.Path(self.get_ext_fullpath(ext.name)) 29 | extdir.parent.mkdir(parents=True, exist_ok=True) 30 | 31 | # example of cmake args 32 | config = 'Debug' if self.debug else 'Release' 33 | cmake_args = [ 34 | '-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_%s=%s' % (config.upper(), str(extdir.parent.absolute())), 35 | '-DCMAKE_BUILD_TYPE=%s' % config 36 | ] 37 | 38 | # example of build args 39 | build_args = [ 40 | '--config', config 41 | ] 42 | 43 | os.chdir(str(build_temp)) 44 | self.spawn(['cmake', str(cwd)] + cmake_args) 45 | if not self.dry_run: 46 | self.spawn(['cmake', '--build', '.'] + build_args) 47 | # Troubleshooting: if fail on line above then delete all possible 48 | # temporary CMake files including "CMakeCache.txt" in top level dir. 49 | os.chdir(str(cwd)) 50 | 51 | ## End from https://stackoverflow.com/questions/42585210/extending-setuptools-extension-to-use-cmake-in-setup-py ## 52 | 53 | # Not put in the setup, but these are the minimum packages 54 | # you should have to use everytihng in the repo. 55 | install_requires=[ 56 | 'pytorch', 57 | 'torch-geometric', 58 | 'torch-scatter', 59 | 'torch-sparse' 60 | 'pytorch-lightning', 61 | 'dotmap', 62 | 'seaborn', 63 | 'numpy', 64 | 'matplotlib' 65 | ] 66 | 67 | setup( 68 | name='automate', 69 | version='1.0.1', 70 | author='Ben Jones', 71 | author_email='benjones@cs.washington.edu', 72 | url='', 73 | description='', 74 | license='MIT', 75 | python_requires='>=3.6', 76 | ext_modules=[CMakeExtension('automate_cpp')], 77 | cmdclass={ 78 | 'build_ext': build_ext 79 | }, 80 | packages=find_packages() 81 | ) --------------------------------------------------------------------------------