add node section overwrite (#5)

* add node section overwrite

* win compat

* just use seek

* forgot ftell

* read and write in binary

* seek to eol after writing line

* turns out seek is only necessary for windows

* should be line_end_pos

src/deck.cpp

@@ -353,22 +353,55 @@ static inline void ans_strtod(char *raw, int fltsz,
 #endif
 }
 
+// FORTRAN-like scientific notation string formatting
+void FormatWithExp(char *buffer, size_t buffer_size, double value, int width,
+                   int precision, int num_exp) {
+  // Format the value using snprintf with given width and precision
+  snprintf(buffer, buffer_size, "% *.*E", width, precision, value);
+
+  // Find the 'E' in the output, which marks the beginning of the exponent.
+  char *exponent_pos = strchr(buffer, 'E');
+
+  // Check the current length of the exponent (after 'E+')
+  int exp_length = strlen(exponent_pos + 2);
+
+  // If the exponent is shorter than the desired number of digits
+  if (exp_length < num_exp) {
+    // Shift the exponent one place to the right and prepend a '0'
+    for (int i = exp_length + 1; i > 0; i--) {
+      exponent_pos[2 + i] = exponent_pos[1 + i];
+    }
+    exponent_pos[2] = '0';
+
+    // Shift the entire buffer to the left to remove the space added by snprintf
+    memmove(
+        buffer, buffer + 1,
+        strlen(
+            buffer)); // length of buffer not recalculated, using previous value
+  }
+}
+
 struct NodeSection {
   NDArray<int, 1> nid;
   NDArray<double, 2> coord;
   NDArray<int, 1> tc;
   NDArray<int, 1> rc;
   int n_nodes = 0;
+  int fpos = 0;
 
   // Default constructor
   NodeSection() {}
 
   NodeSection(std::vector<int> nid_vec, std::vector<double> nodes_vec,
-              std::vector<int> tc_vec, std::vector<int> rc_vec) {
+              std::vector<int> tc_vec, std::vector<int> rc_vec,
+              int file_position) {
     n_nodes = nid_vec.size();
     std::array<int, 1> nid_shape = {static_cast<int>(n_nodes)};
     std::array<int, 2> coord_shape = {static_cast<int>(n_nodes), 3};
 
+    // store file position where node block began
+    fpos = file_position;
+
     // Wrap vectors as NDArrays
     nid = WrapVectorAsNDArray(std::move(nid_vec), nid_shape);
     coord = WrapVectorAsNDArray(std::move(nodes_vec), coord_shape);
@@ -681,6 +714,8 @@ class Deck {
     std::vector<int> rc;
     rc.reserve(NNUM_RESERVE);
 
+    int start_pos = memmap.tellg();
+
     int index = 0;
     while (memmap[0] != '*') {
 
@@ -743,7 +778,7 @@ class Deck {
       memmap.seek_eol();
     }
 
-    NodeSection *node_section = new NodeSection(nid, coord, tc, rc);
+    NodeSection *node_section = new NodeSection(nid, coord, tc, rc, start_pos);
     node_sections.push_back(*node_section);
 
     return;
@@ -862,6 +897,79 @@ class Deck {
   int ReadLine() { return memmap.read_line(); }
 };
 
+void OverwriteNodeSection(const char *filename, int fpos,
+                          const NDArray<const double, 2> coord_arr) {
+
+  // Open the file with read and write permissions
+  FILE *fp = fopen(filename, "rb+");
+  if (!fp) {
+    throw std::runtime_error("Cannot open file for reading and writing.");
+  }
+
+  // Seek to the start position of the node section
+  if (fseek(fp, fpos, SEEK_SET) != 0) {
+    fclose(fp);
+    throw std::runtime_error(
+        "Cannot seek to the start position of the node section.");
+  }
+
+  int node_idx = 0;
+  off_t line_start_pos;
+  char line[256];
+
+  int n_coord = coord_arr.shape(0);
+  const double *coord = coord_arr.data();
+
+  // Prepare buffers for formatted coordinates
+  char coord_str[17] = {0};
+
+  while (fgets(line, sizeof(line), fp)) {
+    // Record the position of the beginning of the line
+    int line_end_pos = ftell(fp);
+    line_start_pos = line_end_pos - strlen(line);
+
+    // Skip comment lines
+    if (line[0] == '$') {
+      continue;
+    }
+
+    // Check for end of node section
+    if (line[0] == '*') {
+      break;
+    }
+
+    // Check if we have more nodes to process
+    if (node_idx >= n_coord) {
+      break;
+    }
+
+    // Ensure the line is long enough
+    size_t line_length = strlen(line);
+    if (line_length < 56) {
+      throw std::runtime_error("Insufficient line length.");
+    }
+
+    // Format the coordinates
+    for (int ii = 0; ii < 3; ii++) {
+      // Generate the formatted string
+      FormatWithExp(coord_str, 17, coord[node_idx * 3 + ii], 16, 9, 2);
+      memcpy(&line[8 + ii * 16], coord_str, 16); // coordinate field
+    }
+
+    // Seek back to the beginning of the line and write the updated line
+    fseek(fp, line_start_pos, SEEK_SET);
+    fwrite(line, 1, line_length, fp);
+#if defined(_WIN32) || defined(_WIN64)
+    fseek(fp, line_end_pos, SEEK_SET);
+#endif
+
+    // Move to the next node
+    node_idx++;
+  }
+
+  fclose(fp);
+}
+
 NB_MODULE(_deck, m) {
   nb::class_<NodeSection>(m, "NodeSection")
       .def(nb::init())
@@ -870,7 +978,8 @@ NB_MODULE(_deck, m) {
       .def_ro("coordinates", &NodeSection::coord, nb::rv_policy::automatic)
       .def_ro("nid", &NodeSection::nid, nb::rv_policy::automatic)
       .def_ro("tc", &NodeSection::tc, nb::rv_policy::automatic)
-      .def_ro("rc", &NodeSection::rc, nb::rv_policy::automatic);
+      .def_ro("rc", &NodeSection::rc, nb::rv_policy::automatic)
+      .def_ro("fpos", &NodeSection::fpos);
 
   nb::class_<ElementSolidSection>(m, "ElementSolidSection")
       .def(nb::init())
@@ -906,4 +1015,6 @@ NB_MODULE(_deck, m) {
       .def("read_element_solid_section", &Deck::ReadElementSolidSection)
       .def("read_element_shell_section", &Deck::ReadElementShellSection)
       .def("read_node_section", &Deck::ReadNodeSection);
+
+  m.def("overwrite_node_section", &OverwriteNodeSection);
 }

src/lsdyna_mesh_reader/_deck.pyi

@@ -21,6 +21,8 @@ class NodeSection:
     @property
     def rc(self) -> IntArray: ...
     def __len__(self) -> int: ...
+    @property
+    def fpos(self) -> int: ...
 
 class ElementSection:
     def __init__(self) -> None: ...
@@ -51,3 +53,5 @@ class _Deck:
     def read_element_shell_section(self) -> None: ...
     def read_node_section(self) -> None: ...
     def read(self) -> None: ...
+
+def overwrite_node_section(filename: str, fpos: int, nodes: FloatArray2D) -> None: ...

src/lsdyna_mesh_reader/deck.py

@@ -1,10 +1,18 @@
+import os
+import shutil
 from pathlib import Path
 from typing import TYPE_CHECKING, List, Union
 
 import numpy as np
 from numpy.typing import NDArray
 
-from lsdyna_mesh_reader._deck import ElementShellSection, ElementSolidSection, NodeSection, _Deck
+from lsdyna_mesh_reader._deck import (
+    ElementShellSection,
+    ElementSolidSection,
+    NodeSection,
+    _Deck,
+    overwrite_node_section,
+)
 
 if TYPE_CHECKING:
     try:
@@ -19,7 +27,7 @@ class Deck:
     Parameters
     ----------
     filename : str | pathlib.Path
-        Path to the keyword file (`*.k`, `*.key`, `*.dyn`).
+        Path to the keyword file (``*.k``, ``*.key``, ``*.dyn``).
 
     Examples
     --------
@@ -35,8 +43,12 @@ class Deck:
 
     def __init__(self, filename: Union[str, Path]) -> None:
         """Initialize the deck object."""
-        self._deck = _Deck(str(filename))
+        filename = str(filename)
+        if not os.path.isfile(filename):
+            raise FileNotFoundError(f"Invalid file or unable to locate {filename}")
+        self._deck = _Deck(filename)
         self._deck.read()
+        self._filename = filename
 
     @property
     def element_solid_sections(self) -> List[ElementSolidSection]:
@@ -235,6 +247,53 @@ def to_grid(self) -> "UnstructuredGrid":
 
         return grid
 
+    def overwrite_node_section(
+        self, filename: Union[str, Path], nodes: NDArray[np.float64]
+    ) -> None:
+        """
+        Create a new deck file with the same content but overwritten node section.
+
+        Parameters
+        ----------
+        filename : str | pathlib.Path
+            Path to the new file.
+        nodes : np.ndarray[float]
+            New node coordinates to write. Number of nodes must match the
+            number of nodes in the node section.
+
+        Notes
+        -----
+        Overwrites the first node section.
+
+        Examples
+        --------
+        Load an existing LS-DYNA deck and write new nodes to it in a new file.
+
+        >>> import numpy as np
+        >>> import lsdyna_mesh_reader
+        >>> from lsdyna_mesh_reader import examples
+        >>> deck = lsdyna_mesh_reader.Deck(examples.birdball)
+        >>> nodes = deck.node_sections[0].coordinates
+        >>> new_nodes = nodes + np.random.random(nodes.shape)
+        >>> deck.overwrite_node_section("new_deck.k", new_nodes)
+
+        """
+        if not self.node_sections:
+            raise RuntimeError("This deck is missing a node section")
+        nsec = self.node_sections[0]
+
+        if nodes.ndim != 2 or nodes.shape[1] != 3:
+            raise ValueError(f"Expected `nodes` to contain 3D coordinates, not {nodes.shape}")
+        elif nodes.shape[0] != nsec.nid.size:
+            raise RuntimeError(
+                f"Number of coordinates in `nodes` ({nodes.shape[0]}) must match number of "
+                f"nodes in the node section ({nsec.nid.size})"
+            )
+
+        filename = str(filename)
+        shutil.copy(self._filename, filename)
+        overwrite_node_section(filename, nsec.fpos, nodes)
+
     def __repr__(self) -> str:
         lines = ["LSDYNA Deck with:"]
         lines.append(f"  Node sections:              {len(self.node_sections)}")

tests/test_deck.py

@@ -203,3 +203,20 @@ def test_element_tshell() -> None:
     assert len(node_section) == 324
 
     grid = deck.to_grid()
+
+
+def test_overwrite_node_section(tmp_path: Path) -> None:
+    filename = str(tmp_path / "tmp.k")
+    with open(filename, "w") as fid:
+        fid.write(NODE_SECTION)
+
+    deck = lsdyna_mesh_reader.Deck(filename)
+
+    new_filename = str(tmp_path / "tmp_overwrite.k")
+    node_section = deck.node_sections[0]
+    new_nodes = node_section.coordinates + np.random.random(node_section.coordinates.shape)
+    deck.overwrite_node_section(new_filename, new_nodes)
+
+    deck_new = lsdyna_mesh_reader.Deck(new_filename)
+    assert np.allclose(deck_new.node_sections[0].coordinates, new_nodes)
+    assert np.allclose(deck_new.node_sections[0].nid, deck.node_sections[0].nid)