Salome HOME
fix on shapeGroups.py
[modules/hydro.git] / src / HYDROTools / shapesGroups.py
1
2 import salome
3 salome.salome_init()
4
5 import  SMESH, SALOMEDS
6 from salome.smesh import smeshBuilder
7
8 import numpy as np
9 import MEDLoader as ml
10 import medcoupling as mc
11
12 import shapefile
13 import math
14 import os
15
16
17 def freeBordersGroup(meshFileIn, meshFileOut=""):
18     """
19     In a mesh, create a group of Edges for the borders: domain, internal holes (isles)
20     parameters:
21     meshFileIn: full path of the input mesh file, format MED
22     meshFileOut: full path of the output mesh file, format MED (default="" : when "", output file is suffixed with "_brd.med"
23     return full path of the output mesh file
24     """
25     smesh = smeshBuilder.New()
26     smesh.SetEnablePublish( False ) # Set to False to avoid publish in study if not needed
27     ([MESH], status) = smesh.CreateMeshesFromMED(meshFileIn)
28
29     nbAdded, MESH, addedBnd = MESH.MakeBoundaryElements( SMESH.BND_1DFROM2D, '', '', 0, [])
30
31     aCriteria = []
32     aCriterion = smesh.GetCriterion(SMESH.EDGE,SMESH.FT_FreeBorders,SMESH.FT_Undefined,0)
33     aCriteria.append(aCriterion)
34     aFilter = smesh.GetFilterFromCriteria(aCriteria)
35     aFilter.SetMesh(MESH.GetMesh())
36     FreeBorders = MESH.GroupOnFilter( SMESH.EDGE, 'FreeBorders', aFilter )
37
38     a = os.path.splitext(meshFileIn)
39     smesh.SetName(MESH, os.path.basename(a[0]))
40
41     if meshFileOut == "":
42         a = os.path.splitext(meshFileIn)
43         smesh.SetName(MESH, os.path.basename(a[0]))
44         meshFileOut = a[0] + '_brd' + a[1]
45
46     try:
47         MESH.ExportMED(meshFileOut,auto_groups=0,minor=40,overwrite=1,meshPart=None,autoDimension=1)
48         pass
49     except:
50         print('ExportMED() failed. Invalid file name?')
51     return meshFileOut
52
53
54 def explodeGroup(grp, grpName):
55     """
56     from a group of edges loaded with MEDCoupling, create ordered lists of nodes, one list for each set of connected edges.
57     parameters:
58     grp: MEDCoupling object for the group of edges
59     grpName: name of the group
60     return:
61     List of descriptors [(ordered list of nodeIds, name of the list, closed status)]
62     """
63     print(" === explodeGroup", grpName)
64     nbCells=grp.getNumberOfCells()
65
66     dicReverse = {} # id node --> id edges
67     for i in range(nbCells):
68         nodcell = grp.getNodeIdsOfCell(i)
69         for j in range(len(nodcell)):
70             if nodcell[j] in dicReverse:
71                 dicReverse[nodcell[j]].append(i)
72             else:
73                 dicReverse[nodcell[j]] = [i]
74
75     nodeChains = []
76     usedCells = [False] * nbCells
77     while False in usedCells:
78         icell = usedCells.index(False)
79         usedCells[icell] = True
80         nodcell = grp.getNodeIdsOfCell(icell)
81         closed = False
82         chain = [nodcell[0], nodcell[1]]
83         nextnode = nodcell[1]
84         prevnode = nodcell[0]
85         while nextnode in dicReverse:
86             nextcells = dicReverse[nextnode]
87             if len(nextcells) != 2:             # end of chain(1) or "edges connector"(>2): stop
88                 closed = False
89                 nextnode = -1                   # stop
90             else:
91                 newcell =False
92                 for i in range(len(nextcells)):
93                     ncell = nextcells[i]
94                     if not usedCells[ncell]:       # the chain of nodes grows
95                         usedCells[ncell] = True
96                         newcell = True
97                         nodcell = grp.getNodeIdsOfCell(ncell)
98                         if nodcell[0] == nextnode:
99                             nextnode = nodcell[1]  # forward edge
100                         else:
101                             nextnode = nodcell[0]  # reversed edge ?
102                         chain.append(nextnode)
103                 if not newcell:                    # end of chain, closed
104                     closed =True
105                     nextnode = -1
106         while prevnode in dicReverse:
107             prevcells = dicReverse[prevnode]
108             if len(prevcells) != 2:             # end of chain(1) or "edges connector"(>2): stop
109                 closed = False
110                 prevnode = -1                   # stop
111             else:
112                 newcell =False
113                 for i in range(len(prevcells)):
114                     ncell = prevcells[i]
115                     if not usedCells[ncell]:       # the chain of nodes grows
116                         usedCells[ncell] = True
117                         newcell = True
118                         nodcell = grp.getNodeIdsOfCell(ncell)
119                         if nodcell[1] == prevnode:
120                             prevnode = nodcell[0]  # forward edge
121                         else:
122                             prevnode = nodcell[1]  # reversed edge ?
123                         chain.insert(0, prevnode)
124                 if not newcell:                    # end of chain, closed
125                     closed =True
126                     prevnode = -1
127
128         chainDesc = (chain, grpName +"_%s" % len(nodeChains), closed)
129         nodeChains.append(chainDesc)
130         print(chainDesc[1:])
131     return nodeChains
132
133
134 def writeShapeLines(mcMesh, grpName, nodeChains, outputDirectory, offsetX=0., offsetY=0.):
135     """
136     from a mesh loaded in memory with MEDLoader, and a list of list of connected nodes associated to a group of edges, write a shapefile of type line, with one record per list of connected nodes.
137     parameters:
138     mcMesh: mesh loaded in memory with MEDLoader
139     grpName: name associated to the group of edges
140     nodeChains: List of descriptors corresponding to the group [(ordered list of nodeIds, name of the list, closed status)]
141     outputDirectory: directory for writing the shapefile
142     offsetX : offset of origin X in the mesh (default 0). The shapefile is always written without local origin to be ready for a direct load in Qgis.
143     offsetY : offset of origin Y in the mesh (default 0). The shapefile is always written without local origin to be ready for a direct load in Qgis.
144     """
145     print(" === writeShapeLines", grpName)
146     coords = mcMesh.getCoords()
147     shapeFileName = os.path.join(outputDirectory, grpName)
148     w = shapefile.Writer(shapeFileName)
149     w.shapeType = 3
150     w.field('name', 'C')
151     for (chain, chainName, closed) in nodeChains:
152         print("   --- ", chainName)
153         chaincoords = []
154         for node in chain:
155             coord = coords[node].getValues()
156             coordLb93=[coord[0] + offsetX, coord[1] + offsetY]
157             chaincoords.append(coordLb93)
158         w.line([chaincoords])
159         w.record(chainName)
160     w.close()
161
162
163 def writeShapePoints(mcMesh, grpName, nodeChains, outputDirectory, offsetX=0., offsetY=0.):
164     """
165     from a mesh loaded in memory with MEDLoader, and a list of list of connected nodes associated to a group of edges, write a shapefile of type multi points, with one record per list of connected nodes.
166     parameters:
167     mcMesh: mesh loaded in memory with MEDLoader
168     grpName: name associated to the group of edges
169     nodeChains: List of descriptors corresponding to the group [(ordered list of nodeIds, name of the list, closed status)]
170     outputDirectory: directory for writing the shapefile
171     offsetX : offset of origin X in the mesh (default 0). The shapefile is always written without local origin to be ready for a direct load in Qgis.
172     offsetY : offset of origin Y in the mesh (default 0). The shapefile is always written without local origin to be ready for a direct load in Qgis.
173     """
174     print(" === writeShapePoints", grpName)
175     coords = mcMesh.getCoords()
176     shapeFileName = os.path.join(outputDirectory, grpName)
177     w = shapefile.Writer(shapeFileName + '_pts')
178     #w.shapeType = 8
179     w.field('name', 'C')
180     for (chain, chainName, closed) in nodeChains:
181         print("   --- ", chainName)
182         chaincoords = []
183         for node in chain:
184             coord = coords[node].getValues()
185             coordLb93=[coord[0] + offsetX, coord[1] + offsetY]
186             chaincoords.append(coordLb93)
187         w.multipoint(chaincoords)
188         w.record(chainName)
189     w.close()
190
191
192 def exploreEdgeGroups(meshFile, outputDirectory="", offsetX=0., offsetY=0.):
193     """
194     Find all the groups of edges in a mesh and, for each group, create one shapefile of lines and one of points. The shapefiles are created in the same system of coordinates as the mesh (For instance Lambert 93), but without origin offset (to be ready for a direct load in Qgis)
195     parameters:
196     meshFile: full path of the input mesh file, format MED
197     outputDirectory: directory in which the shapefiles are written (default "", if "" use the directory containing the mesh
198     offsetX: local X origin of the mesh
199     offsetY: local Y origin of the mesh
200     """
201     print(" === exploreEdgeGroups", meshFile)
202     if outputDirectory == "":
203         outputDirectory = os.path.dirname(meshFile)
204
205     a = os.path.splitext(meshFile)
206     prefix = os.path.basename(a[0]) # prefix = file name without extension
207
208     mcMesh = ml.MEDFileMesh.New(meshFile)
209     dim = mcMesh.getMeshDimension()
210     d1=-1        # when dimension 2, edges are dim -1
211     if dim == 3: # when dimension 3, edges are dim -2
212         d1=-2
213
214     grp_names = mcMesh.getGroupsOnSpecifiedLev(d1) #names of edges groups
215
216     groups = [mcMesh.getGroup(d1, name) for name in grp_names] # list of groups in their name order
217
218     for (grp, grpName) in zip(groups, grp_names):
219         fullGrpName = prefix + '_' + grpName
220         nodeChains = explodeGroup(grp, fullGrpName)
221         writeShapeLines(mcMesh, fullGrpName, nodeChains, outputDirectory, offsetX, offsetY)
222         writeShapePoints(mcMesh, fullGrpName, nodeChains, outputDirectory, offsetX, offsetY)
223
224
225 def fitShapePointsToMesh(freeBorderShapefile, shapefileToAdjust, splitFreeBorder=True, splitShapeAdjusted=True):
226     """
227     shapeFileToAdjust must be a closed line or polygon crossing freeBorderShapefile in 2 points.
228     Find in shapeFileToAdjust and in freeBorderShapefile the two closest corresponding points and move the points in shapeFileToAdjust to correspond to the points found in freeBorderShapefile. Split shapeFileToAdjust in two parts (inside or outside freeBorder). If requested, split freeBorderShapefile in two parts. Same for shapeFileToAdjust.
229     parameters:
230     freeBorderShapefile: a set of free border lines, as generated by the functions freeBordersGroup and exploreEdgeGroups.
231     shapefileToAdjust: a closed line or polygon, supposed to be drawn in qgis to pass as close as possible to the points to be connected, on the free border.
232     splitFreeBorder: boolean default True
233     splitShapeAdjusted: boolean default True
234     """
235     print(" === fitShapePointsToMesh", freeBorderShapefile, shapefileToAdjust)
236
237     # --- find domain freeBorder:  bounding box englobing all others
238     #     TODO: This may not be always the case, when there is not a single domain with holes, but several non connected domains.
239
240     fb = shapefile.Reader(freeBorderShapefile)
241     fbShapes = fb.shapes()
242     maxbbox=[1.e30, 1.e30, -1.e30, -1.e30]
243     outerBboxIndex = -1
244     for i,s in enumerate(fbShapes):
245         bb = s.bbox
246         if (bb[0] < maxbbox[0] and bb[1] < maxbbox[1] and bb[2] > maxbbox[2] and bb[3] > maxbbox[3]):
247             maxbbox = bb
248             outerBboxIndex = i
249     fbs = fbShapes[outerBboxIndex] # the domain free border shape
250
251     # --- find the intersections of the shapefile to adjust and the domain free border:
252     #     the closests points (two pairs of points)
253
254     sf = shapefile.Reader(shapefileToAdjust)
255     shapes = sf.shapes()
256     sfta = sf.shape(0)
257     pdist =[]
258     x0 = -1.e30
259     y0 = -1.e30
260     for i,p in enumerate(sfta.points):
261         d = math.sqrt((x0-p[0])*(x0-p[0]) + (y0-p[1])*(y0-p[1]))
262         x0 = p[0]
263         y0 = p[1]
264         if d < 1.e-5:
265             continue # do not take into account superposed points in shape to adjust
266         dmin = 1.e30
267         jmin = -1
268         for j,pfb in enumerate(fbs.points):
269             d = math.sqrt((pfb[0]-p[0])*(pfb[0]-p[0]) + (pfb[1]-p[1])*(pfb[1]-p[1]))
270             if d < dmin:
271                 dmin = d
272                 jmin = j
273         pdist.append((dmin, jmin, i)) # distance, index in freeBorder, index in shapeToAdjust
274     pdist.sort() # the 2 closest points must be set on the mesh freeBorder
275     print(pdist)
276     i1 = min(pdist[0][2], pdist[1][2]) # index of first adjusted point in shapeToAdjust
277     i2 = max(pdist[0][2], pdist[1][2]) # index of second adjusted point in shapeToAdjust
278     if i1 == pdist[0][2]:
279         ifb1 = pdist[0][1]             # index of first adjusted point in freeBorder
280         ifb2 = pdist[1][1]             # index of second adjusted point in freeBorder
281     else:
282         ifb1 = pdist[1][1]
283         ifb2 = pdist[0][1]
284     print("i1, i2, len(sfta.points)", i1, i2, len(sfta.points))
285     print("ifb1, ifb2, len(fbs.points)", ifb1, ifb2, len(fbs.points))
286
287     # --- write the adusted shapefile: free border closest points replaced with corresponding points
288     #     on the free border. two polylines, one inside the domain, one outside
289
290     a = os.path.splitext(shapefileToAdjust)
291     shapefileAdjusted = a[0] + '_adj' + a[1]
292     chainName = os.path.basename(a[0])
293
294     w = shapefile.Writer(shapefileAdjusted)
295     w.shapeType = 3
296     w.field('name', 'C')
297
298     if splitShapeAdjusted:
299         chaincoords = []
300         chaincoords.append(fbs.points[ifb1])
301         for i in range(i1+1, i2):
302             chaincoords.append(sfta.points[i])
303         chaincoords.append(fbs.points[ifb2])
304         w.line([chaincoords])
305         w.record(chainName + '_0')
306
307         chaincoords = []
308         chaincoords.append(fbs.points[ifb2])
309         if i2+1 < len(sfta.points):
310             for i in range(i2+1, len(sfta.points)):
311                 chaincoords.append(sfta.points[i])
312         for i in range(i1):
313             chaincoords.append(sfta.points[i])
314         chaincoords.append(fbs.points[ifb1])
315         w.line([chaincoords])
316         w.record(chainName + '_1')
317     else:
318         chaincoords = []
319         chaincoords.append(fbs.points[ifb1])
320         for i in range(i1+1, i2):
321             chaincoords.append(sfta.points[i])
322         chaincoords.append(fbs.points[ifb2])
323         if i2+1 < len(sfta.points):
324             for i in range(i2+1, len(sfta.points)):
325                 chaincoords.append(sfta.points[i])
326         for i in range(i1):
327             chaincoords.append(sfta.points[i])
328         w.line([chaincoords])
329         w.record(chainName)
330
331     w.close()
332
333     if splitFreeBorder:
334         # write the free border split in two polylines (cut by the adjusted shapefile)
335
336         a = os.path.splitext(freeBorderShapefile)
337         freeBorderSplit = a[0] + '_split' + a[1]
338         chainName = os.path.basename(a[0])
339
340         w = shapefile.Writer(freeBorderSplit)
341         w.shapeType = 3
342         w.field('name', 'C')
343
344         if (ifb1 > ifb2):
345             i = ifb1; ifb1 = ifb2; ifb2 = i
346
347         chaincoords = []
348         for i in range(ifb1, ifb2+1):
349             chaincoords.append(fbs.points[i])
350         w.line([chaincoords])
351         w.record(chainName + '_0')
352
353         chaincoords = []
354         for i in range(ifb2, len(fbs.points)):
355             chaincoords.append(fbs.points[i])
356         for i in range(ifb1+1):
357             chaincoords.append(fbs.points[i])
358         w.line([chaincoords])
359         w.record(chainName + '_1')
360         w.close()
361