combinatorics.py 5,8 КБ
Newer Older
DolphinDream's avatar
DolphinDream включено в состав коммита
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
# ##### BEGIN GPL LICENSE BLOCK #####
#
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU General Public License
#  as published by the Free Software Foundation; either version 2
#  of the License, or (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software Foundation,
#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####

import bpy
DolphinDream's avatar
DolphinDream включено в состав коммита
20
from bpy.props import IntProperty, EnumProperty
DolphinDream's avatar
DolphinDream включено в состав коммита
21
22
23
24

from sverchok.node_tree import SverchCustomTreeNode
from sverchok.data_structure import (match_long_repeat, updateNode)

DolphinDream's avatar
DolphinDream включено в состав коммита
25
from itertools import (product, permutations, combinations)
DolphinDream's avatar
DolphinDream включено в состав коммита
26

DolphinDream's avatar
DolphinDream включено в состав коммита
27
28
29
30
operations = {
    "PRODUCT":      (10, lambda s, r: product(*s, repeat=r)),
    "PERMUTATIONS": (20, lambda s, l: permutations(s, l)),
    "COMBINATIONS": (30, lambda s, l: combinations(s, l))
DolphinDream's avatar
DolphinDream включено в состав коммита
31
32
}

DolphinDream's avatar
DolphinDream включено в состав коммита
33
34
35
36
37
operationItems = [(k, k.title(), "", s[0]) for k, s in sorted(operations.items(), key=lambda k: k[1][0])]

ABC = tuple('ABCDEFGHIJKLMNOPQRSTUVWXYZ')  # input socket labels

multiple_input_operations = {"PRODUCT"}
DolphinDream's avatar
DolphinDream включено в состав коммита
38
39
40
41


class SvCombinatoricsNode(bpy.types.Node, SverchCustomTreeNode):
    """
DolphinDream's avatar
DolphinDream включено в состав коммита
42
    Triggers: Product, Permutations, Combinations
DolphinDream's avatar
DolphinDream включено в состав коммита
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
    Tooltip: Generate various combinatoric operations
    """
    bl_idname = 'SvCombinatoricsNode'
    bl_label = 'Combinatorics'

    def update_operation(self, context):
        self.label = self.operation.title()
        self.update_sockets()
        updateNode(self, context)

    operation = EnumProperty(
        name="Operation", items=operationItems,
        description="Operation type", default="PRODUCT",
        update=update_operation)

    repeat = IntProperty(
DolphinDream's avatar
DolphinDream включено в состав коммита
59
        name='Repeat', description='Repeat the list inputs this many times',
DolphinDream's avatar
DolphinDream включено в состав коммита
60
61
62
        default=1, min=1, update=updateNode)

    length = IntProperty(
DolphinDream's avatar
DolphinDream включено в состав коммита
63
64
        name='Length', description='Limit the elements to operate on to this value',
        default=1, min=0, update=updateNode)
DolphinDream's avatar
DolphinDream включено в состав коммита
65
66
67
68
69
70
71
72
73

    def sv_init(self, context):
        self.inputs.new('StringsSocket', "Repeat").prop_name = "repeat"
        self.inputs.new('StringsSocket', "Length").prop_name = "length"
        self.inputs.new('StringsSocket', "A")
        self.inputs.new('StringsSocket', "B")

        self.outputs.new('StringsSocket', "Result")

DolphinDream's avatar
DolphinDream включено в состав коммита
74
        self.update_operation(context)
DolphinDream's avatar
DolphinDream включено в состав коммита
75
76
77
78

    def update(self):
        ''' Add/remove sockets as A-Z sockets are connected/disconnected '''

DolphinDream's avatar
DolphinDream включено в состав коммита
79
        # not a multiple input operation ? => no need to update sockets
DolphinDream's avatar
DolphinDream включено в состав коммита
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
        if self.operation not in multiple_input_operations:
            return

        inputs = self.inputs

        # get all existing A-Z sockets (connected or not)
        inputs_AZ = list(filter(lambda s: s.name in ABC, inputs))

        # last A-Z socket connected ? => add an empty A-Z socket at the end
        if inputs_AZ[-1].links:
            name = ABC[len(inputs_AZ)]  # pick the next letter A to Z
            inputs.new("StringsSocket", name)

        else:  # last input disconnected ? => remove all but last unconnected
            while len(inputs_AZ) > 2 and not inputs_AZ[-2].links:
                s = inputs_AZ[-1]
                inputs.remove(s)
                inputs_AZ.remove(s)

    def update_sockets(self):
DolphinDream's avatar
DolphinDream включено в состав коммита
100
        ''' Update sockets based on selected operation '''
DolphinDream's avatar
DolphinDream включено в состав коммита
101
102
103
104
105
106
107
108
109
110
111
112

        inputs = self.inputs

        # update the A-Z input sockets
        if self.operation in multiple_input_operations:
            if not "B" in inputs:
                inputs.new("StringsSocket", "B")
        else:
            for a in ABC[1:]:  # remove all B-Z inputs (keep A)
                if a in inputs:
                    inputs.remove(inputs[a])

DolphinDream's avatar
DolphinDream включено в состав коммита
113
114
115
116
        # update the other sockets
        if self.operation in {"PRODUCT"}:
            if inputs["Repeat"].hide:
                inputs["Repeat"].hide_safe = False
DolphinDream's avatar
DolphinDream включено в состав коммита
117
            inputs["Length"].hide_safe = True
DolphinDream's avatar
DolphinDream включено в состав коммита
118
        elif self.operation in {"COMBINATIONS", "PERMUTATIONS"}:
DolphinDream's avatar
DolphinDream включено в состав коммита
119
            inputs["Repeat"].hide_safe = True
DolphinDream's avatar
DolphinDream включено в состав коммита
120
121
            if inputs["Length"].hide:
                inputs["Length"].hide_safe = False
DolphinDream's avatar
DolphinDream включено в состав коммита
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137

    def draw_buttons(self, context, layout):
        layout.prop(self, "operation", text="")

    def process(self):
        outputs = self.outputs
        # return if no outputs are connected
        if not any(s.is_linked for s in outputs):
            return

        inputs = self.inputs

        all_AZ_sockets = list(filter(lambda s: s.name in ABC, inputs))
        connected_AZ_sockets = list(filter(lambda s: s.is_linked, all_AZ_sockets))

        # collect the data inputs from all connected AZ sockets
DolphinDream's avatar
DolphinDream включено в состав коммита
138
        I = [s.sv_get()[0] for s in connected_AZ_sockets]
DolphinDream's avatar
DolphinDream включено в состав коммита
139
140
141
142
143

        if self.operation == "PRODUCT":
            R = inputs["Repeat"].sv_get()[0]
            R = list(map(lambda x: max(1, int(x)), R))
            parameters = match_long_repeat([[I], R])
DolphinDream's avatar
DolphinDream включено в состав коммита
144
        else:  # PERMUTATIONS / COMBINATIONS
DolphinDream's avatar
DolphinDream включено в состав коммита
145
            L = inputs["Length"].sv_get()[0]
DolphinDream's avatar
DolphinDream включено в состав коммита
146
            L = list(map(lambda x: max(0, int(x)), L))
DolphinDream's avatar
DolphinDream включено в состав коммита
147
148
            parameters = match_long_repeat([I, L])

DolphinDream's avatar
DolphinDream включено в состав коммита
149
150
151
152
153
154
155
156
157
        function = operations[self.operation][1]

        resultList = []
        for sequence, v in zip(*parameters):
            if self.operation in {"PERMUTATIONS", "COMBINATIONS"}:
                if v == 0 or v > len(sequence):
                    v = len(sequence)
            result = [list(a) for a in function(sequence, v)]
            resultList.append(result)
DolphinDream's avatar
DolphinDream включено в состав коммита
158
159
160
161
162
163
164
165
166
167

        outputs["Result"].sv_set(resultList)


def register():
    bpy.utils.register_class(SvCombinatoricsNode)


def unregister():
    bpy.utils.unregister_class(SvCombinatoricsNode)