### ===================================================================================================================
### FUNCTION: Add VIIA model picture legend
### ===================================================================================================================
# Copyright ©VIIA 2024
### ===================================================================================================================
### 1. Import modules
### ===================================================================================================================
# General imports
from __future__ import annotations
import json
from pathlib import Path
from typing import TYPE_CHECKING, List, Dict, Optional, Union
# References for functions and classes in the rhdhv_fem package
from rhdhv_fem.fem_config import Config
from rhdhv_fem.materials import Concrete, Timber, DirectStiffnessMatrixModel
from rhdhv_fem.shapes import Shapes, Lintel
# References for functions and classes in the viiaPackage
if TYPE_CHECKING:
from viiapackage.viiaStatus import ViiaProject
### ===================================================================================================================
### 2. Functions to load legend data and add legends to model pictures
### ===================================================================================================================
[docs]def viia_model_picture_legend_data(
project: ViiaProject, model_picture_name: str, json_dir: Path, shape_colours: List[str], shapes: List[Shapes],
element: str = None) -> Optional[Union[Dict[str, str], Dict[str, Dict[str, str]]]]:
"""
This function will write a json file to output the project pictures information. Each picture is output with all
color codes in that picture and a name associated with each color code.
Input:
- project (obj): Project object containing collections and of fem objects and project variables.
- model_picture_name (str): The name of the model picture to add the legend data.
- json_dir (Path): Path that json file will be created in.
- shapes_colours (list of str): The list contains the colour codes included in the model picture.
- shapes (list of shapes): The list contains the shapes included in the model picture.
- element (str): The string indicates the element category, e.g., beam.
Output:
- A json file containing all picture colour codes and names will be created.
"""
def _get_mat_name(material_name, name_map):
""" Get the material name given a dictionary mapping, if it is not in the dictionary, returns other material."""
for k, v in name_map.items():
if k in material_name:
mat_name = v
break
else:
mat_name = 'Other material'
return mat_name
if model_picture_name and element:
# Gather the colour dictionary
if element == 'column':
colour_dictionary = getattr(project, f'beam_colour_dictionary', {})
else:
colour_dictionary = getattr(project, f'{element}_colour_dictionary', {})
if not colour_dictionary:
project.write_log(
f"WARNING: The colour dictionary for {element} is not found, the legend info is not be generated.")
return None
else:
project.write_log(
"WARNING: Model picture name or element category is not given, the legend info is not be generated.")
return None
# Get the unique colours
unique_shape_colours = list(set(shape_colours))
if '#ffffff' in shape_colours:
unique_shape_colours.remove('#ffffff')
# Gather the colours and legends
legend_data = {}
colour_shape_dict = {colour: shapes[shape_colours.index(colour)] for colour in unique_shape_colours}
if not all([hasattr(shape, 'geometry') and hasattr(shape, 'material') for shape in colour_shape_dict.values()]):
project.write_log(
f"WATNING: Not all shapes in model picture {model_picture_name} have material and geomtry, the legend "
f"info cannot be generated.")
return legend_data
# Store the correct legend name with the colour as the key in the dictionary
if element == 'wall':
for colour, shape in colour_shape_dict.items():
if isinstance(shape, Lintel) and 'LATEI' in shape.name:
name_get = _get_mat_name(shape.material.name,
project.project_specific['name_mapping']['name_map_beams'])
geom_name = shape.geometry.name.replace('LATEI', 'lintel')
legend_data[colour] = f'{name_get} {geom_name}'
else:
name_get = _get_mat_name(shape.material.name,
project.project_specific['name_mapping']['name_map_walls'])
legend_data[colour] = f'{name_get} {shape.geometry.geometry_model.thickness * 1000}mm'
elif element in ['roof', 'floor', 'fstrip']:
for colour, shape in colour_shape_dict.items():
name_get = _get_mat_name(shape.material.name, project.project_specific['name_mapping']['name_map_floors'])
if isinstance(shape.material.material_model, DirectStiffnessMatrixModel):
legend_data[colour] = f'{name_get} {shape.geometry.geometry_model.thickness * 1000}mm (equivalent)'
else:
legend_data[colour] = f'{name_get} {shape.geometry.geometry_model.thickness * 1000}mm'
elif element == 'beam':
for colour, shape in colour_shape_dict.items():
name_get = _get_mat_name(shape.material.name, project.project_specific['name_mapping']['name_map_beams'])
geom_name = shape.geometry.name.replace('BALK', 'beam')
legend_data[colour] = f'{name_get} {geom_name}'
elif element == 'column':
for colour, shape in colour_shape_dict.items():
name_get = _get_mat_name(shape.material.name, project.project_specific['name_mapping']['name_map_beams'])
geom_name = shape.geometry.name.replace('KOLOM', 'column')
legend_data[colour] = f'{name_get} {geom_name}'
elif element == 'pile':
for colour, shape in colour_shape_dict.items():
if isinstance(shape.material, Concrete):
legend_data[colour] = f'Concreate pile {shape.geometry.geometry_model.name} mm'
elif isinstance(shape.material, Timber):
legend_data[colour] = f'Timber pile {shape.geometry.geometry_model.name} mm'
else:
legend_data[colour] = f'pile {shape.geometry.geometry_model.name} mm'
model_pic_legend = {model_picture_name: legend_data}
return model_pic_legend
[docs]def viia_add_legend_model_pictures(
project: ViiaProject, json_dir: Path, output_dir: Optional[Path] = None, legend_loc: str = 'upper right',
fontsize: int = 7) -> None:
"""
This function will add a legend to the existing figures in the model picture folder according to the legend_data
json-file in the folder.
Input:
- project (obj): Project object containing collections and of fem objects and project variables.
- json_dir (Path): Path that contains the legend_data.json and model pictures.
- output_dir (Path): Path that output the model pictures with legends. Default value None, using the work
folder to store the pictures.
- legend_loc (str): The legend location which is used for matplotlib. Default location is upper right.
- fontsize (int): The font size of the legend used for matplotlib. Default value 7.
Output:
- A json-file containing all picture color codes and names will be created.
"""
import matplotlib
matplotlib.use(Config.MPL_NONINTERACTIVE(notify=False))
import matplotlib.pyplot as plt
import matplotlib.cbook as cbook
import matplotlib.image as image
from matplotlib.lines import Line2D
import shutil
json_name = json_dir / 'legend_info.json'
# Check if the file is there
if not json_name.exists():
project.write_log(
f"ERROR: The legend_info json-file is not in folder {json_dir.as_posix()}, please provide the file if you "
f"want a legend in the pictures.")
return
# To keep the original figures, send a copy
new_path = json_dir / 'original_figures'
new_path.mkdir(exist_ok=True)
for item in json_dir.iterdir():
if item.is_file():
shutil.copy(str(item), str(new_path / item.name))
project.write_log(f"The original pictures have been saved to {new_path}.")
save_folder = json_dir
if output_dir:
save_folder = output_dir
with json_name.open('r') as f:
output = json.load(f)
# Open figure name
for k, v in output.items():
# Add legend if the legend information in not empty
if v:
with cbook.get_sample_data(json_dir / f'{k}.png') as file:
im = image.imread(file)
fig, ax = plt.subplots(figsize=(12, 6), dpi=300)
# Create legend instance in matplotlib
legend_elements = [Line2D([0], [0], color=f'{i}', lw=2, label=f'{v[i]}') for i in v]
# Create the figure
ax.imshow(im)
ax.legend(handles=legend_elements, loc=legend_loc, fontsize=fontsize)
ax.set_xticks([])
ax.set_yticks([])
fig.patch.set_visible(False)
ax.axis('off')
# Output figure with legend
fig.savefig(f'{save_folder / k}.png')
project.write_log(f'Legend has been added to {save_folder / k}.png.')
else:
project.write_log(
f"WARNING: The legend information for {save_folder / k}.png is not found, please check.")
project.write_log("Adding the legend is done.")
### ===================================================================================================================
### 3. End of script
### ===================================================================================================================