1/*******************************************************************************
2 Copyright (c) The Taichi Authors (2016- ). All Rights Reserved.
3 The use of this software is governed by the LICENSE file.
4*******************************************************************************/
5
6#pragma once
7
8#include <cstring>
9#include <cstdio>
10#include <string>
11#include <vector>
12#include <iterator>
13
14#include "array_fwd.h"
15#include "linalg.h"
16
17namespace taichi {
18
19template <>
20class IndexND<2> {
21 private:
22 int x_[2], y_[2];
23
24 public:
25 using Index = IndexND<2>;
26
27 int i, j;
28 // int offset;
29 int stride;
30 Vector2 storage_offset;
31
32 IndexND() {
33 }
34
35 IndexND(int x0,
36 int x1,
37 int y0,
38 int y1,
39 Vector2 storage_offset = Vector2(0.5f, 0.5f)) {
40 x_[0] = x0;
41 x_[1] = x1;
42 y_[0] = y0;
43 y_[1] = y1;
44 i = x_[0];
45 j = y_[0];
46 // offset = 0;
47 stride = y_[1] - y_[0];
48 this->storage_offset = storage_offset;
49 }
50
51 IndexND(Vector2i start,
52 Vector2i end,
53 Vector2 storage_offset = Vector2(0.5f, 0.5f)) {
54 x_[0] = start[0];
55 x_[1] = end[0];
56 y_[0] = start[1];
57 y_[1] = end[1];
58 i = x_[0];
59 j = y_[0];
60 // offset = 0;
61 stride = y_[1] - y_[0];
62 this->storage_offset = storage_offset;
63 }
64
65 IndexND(int i, int j) {
66 this->i = i;
67 this->j = j;
68 }
69
70 void next() {
71 j++;
72 // offset++;
73 if (j == y_[1]) {
74 j = y_[0];
75 i++;
76 if (i == x_[1]) {
77 }
78 }
79 }
80
81 Index operator++() {
82 this->next();
83 return *this;
84 }
85
86 bool operator==(const IndexND<2> &o) const {
87 return (i == o.i && j == o.j);
88 }
89
90 bool operator!=(const IndexND<2> &o) const {
91 return !(i == o.i && j == o.j);
92 }
93
94 Index &to_end() {
95 i = x_[1];
96 j = y_[0];
97 // offset = (x[1] - x[0]) * (y[1] - y[0]);
98 return *this;
99 }
100
101 const Index &operator*() const {
102 return *this;
103 }
104
105 Index &operator*() {
106 return *this;
107 }
108
109 int operator[](int c) {
110 return *(&i + c);
111 }
112
113 int operator[](int c) const {
114 return *(&i + c);
115 }
116
117 Index neighbour(int di, int dj) const {
118 Index i = *this;
119 i.i += di;
120 i.j += dj;
121 return i;
122 }
123
124 Index neighbour(Vector2i d) const {
125 Index i = *this;
126 i.i += d.x;
127 i.j += d.y;
128 return i;
129 }
130
131 Index operator+(Vector2i d) const {
132 return neighbour(d);
133 }
134
135 Vector2 get_pos() const {
136 return Vector2((real)i + storage_offset.x, (real)j + storage_offset.y);
137 }
138
139 Vector2i get_ipos() const {
140 return Vector2i(i, j);
141 }
142};
143
144typedef IndexND<2> Index2D;
145
146template <>
147class RegionND<2> {
148 private:
149 int x_[2], y_[2];
150 Index2D index_begin_;
151 Index2D index_end_;
152 Vector2 storage_offset_;
153
154 public:
155 using Region = RegionND<2>;
156
157 RegionND() {
158 }
159
160 RegionND(int x0,
161 int x1,
162 int y0,
163 int y1,
164 Vector2 storage_offset = Vector2(0.5f, 0.5f)) {
165 x_[0] = x0;
166 x_[1] = x1;
167 y_[0] = y0;
168 y_[1] = y1;
169 index_begin_ = Index2D(x0, x1, y0, y1, storage_offset);
170 index_end_ = Index2D(x0, x1, y0, y1, storage_offset).to_end();
171 this->storage_offset_ = storage_offset;
172 }
173
174 RegionND(Vector2i start,
175 Vector2i end,
176 Vector2 storage_offset = Vector2(0.5f, 0.5f)) {
177 x_[0] = start[0];
178 x_[1] = end[0];
179 y_[0] = start[1];
180 y_[1] = end[1];
181 index_begin_ = Index2D(start, end, storage_offset);
182 index_end_ = Index2D(start, end, storage_offset).to_end();
183 this->storage_offset_ = storage_offset;
184 }
185
186 const Index2D begin() const {
187 return index_begin_;
188 }
189
190 Index2D begin() {
191 return index_begin_;
192 }
193
194 const Index2D end() const {
195 return index_end_;
196 }
197
198 Index2D end() {
199 return index_end_;
200 }
201};
202
203typedef RegionND<2> Region2D;
204
205template <typename T>
206class ArrayND<2, T> {
207 protected:
208 Region2D region;
209 typedef typename std::vector<T>::iterator iterator;
210 int size;
211 Vector2i res;
212 Vector2 storage_offset = Vector2(0.5f, 0.5f); // default : center storage
213 public:
214 std::vector<T> data;
215 template <typename S>
216 using Array2D = ArrayND<2, S>;
217
218 template <typename P>
219 friend Array2D<T> operator*(const P &b, const Array2D<T> &a);
220
221 int get_size() const {
222 return size;
223 }
224
225 const Region2D &get_region() const {
226 return region;
227 }
228
229 explicit ArrayND(const Vector2i &res,
230 T init = T(0),
231 Vector2 storage_offset = Vector2(0.5f)) {
232 initialize(res, init, storage_offset);
233 }
234
235 void initialize(const Vector2i &res,
236 T init = T(0),
237 Vector2 storage_offset = Vector2(0.5f)) {
238 this->res = res;
239 region = Region2D(0, res[0], 0, res[1], storage_offset);
240 size = res[0] * res[1];
241 data = std::vector<T>(size, init);
242 this->storage_offset = storage_offset;
243 }
244
245 Array2D<T> same_shape(T init) const {
246 return ArrayND<2, T>(res, init, storage_offset);
247 }
248
249 Array2D<T> same_shape() const {
250 return ArrayND<2, T>(res);
251 }
252
253 ArrayND(const Array2D<T> &arr) : ArrayND(arr.res) {
254 this->data = arr.data;
255 this->storage_offset = arr.storage_offset;
256 }
257
258 template <typename P>
259 Array2D<T> operator*(const P &b) const {
260 Array2D<T> o(res);
261 for (int i = 0; i < size; i++) {
262 o.data[i] = b * data[i];
263 }
264 return o;
265 }
266
267 template <typename P>
268 Array2D<T> operator/(const P &b) const {
269 b = T(1) / b;
270 return b * (*this);
271 }
272
273 Array2D<T> operator+(const Array2D<T> &b) const {
274 Array2D<T> o(res);
275 assert(same_dim(b));
276 for (int i = 0; i < size; i++) {
277 o.data[i] = data[i] + b.data[i];
278 }
279 return o;
280 }
281
282 Array2D<T> operator-(const Array2D<T> &b) const {
283 Array2D<T> o(res);
284 assert(same_dim(b));
285 for (int i = 0; i < size; i++) {
286 o.data[i] = data[i] - b.data[i];
287 }
288 return o;
289 }
290
291 void operator+=(const Array2D<T> &b) {
292 assert(same_dim(b));
293 for (int i = 0; i < size; i++) {
294 data[i] = data[i] + b.data[i];
295 }
296 }
297
298 void operator-=(const Array2D<T> &b) {
299 assert(same_dim(b));
300 for (int i = 0; i < size; i++) {
301 data[i] = data[i] - b.data[i];
302 }
303 }
304
305 Array2D<T> &operator=(const Array2D<T> &arr) {
306 this->res = arr.res;
307 this->size = arr.size;
308 this->data = arr.data;
309 this->region = arr.region;
310 this->storage_offset = arr.storage_offset;
311 return *this;
312 }
313
314 Array2D<T> &operator=(const T &a) {
315 for (int i = 0; i < size; i++) {
316 data[i] = a;
317 }
318 return *this;
319 }
320
321 ArrayND() {
322 res = Vector2i(0);
323 size = 0;
324 data.resize(0);
325 }
326
327 ~ArrayND() {
328 }
329
330 void reset(T a) {
331 for (int i = 0; i < size; i++) {
332 data[i] = a;
333 }
334 }
335
336 void reset_zero() {
337 memset(&data[0], 0, sizeof(T) * data.size());
338 }
339
340 bool same_dim(const Array2D<T> &arr) const {
341 return res == arr.res;
342 }
343
344 T dot(const Array2D<T> &b) const {
345 T sum = 0;
346 assert(same_dim(b));
347 for (int i = 0; i < size; i++) {
348 sum += this->data[i] * b.data[i];
349 }
350 return sum;
351 }
352
353 double dot_double(const Array2D<T> &b) const {
354 double sum = 0;
355 assert(same_dim(b));
356 for (int i = 0; i < size; i++) {
357 sum += this->data[i] * b.data[i];
358 }
359 return sum;
360 }
361
362 Array2D<T> add(T alpha, const Array2D<T> &b) const {
363 Array2D<T> o(res);
364 assert(same_dim(b));
365 for (int i = 0; i < size; i++) {
366 o.data[i] = data[i] + alpha * b.data[i];
367 }
368 return o;
369 }
370
371 void add_in_place(T alpha, const Array2D<T> &b) {
372 for (int i = 0; i < size; i++) {
373 data[i] += alpha * b.data[i];
374 }
375 }
376
377 T *operator[](int i) {
378 return &data[0] + i * res[1];
379 }
380
381 const T *operator[](int i) const {
382 return &data[0] + i * res[1];
383 }
384
385 const T &get(int i, int j) const {
386 return (*this)[i][j];
387 }
388
389 const T &get(const Index2D &ind) const {
390 return get(ind.i, ind.j);
391 }
392
393 T get_copy(int i, int j) const {
394 return (*this)[i][j];
395 }
396
397 void set(int i, int j, const T &t) {
398 (*this)[i][j] = t;
399 }
400
401 void set(const Index2D &ind, const T &t) {
402 (*this)[ind] = t;
403 }
404
405 T abs_sum() const {
406 T ret = 0;
407 for (int i = 0; i < size; i++) {
408 ret += std::abs(data[i]);
409 }
410 return ret;
411 }
412
413 T sum() const {
414 T ret = 0;
415 for (int i = 0; i < size; i++) {
416 ret += data[i];
417 }
418 return ret;
419 }
420
421 template <typename TT = T,
422 typename std::enable_if_t<!std::is_class<TT>::value, int> = 0>
423 T abs_max() const {
424 T ret(0);
425 for (int i = 0; i < size; i++) {
426 ret = std::max(ret, abs(data[i]));
427 }
428 return ret;
429 }
430
431 template <typename TT = T, typename TS = typename TT::ScalarType>
432 TS abs_max() const {
433 TS ret(0);
434 for (int i = 0; i < size; i++) {
435 ret = std::max(ret, data[i].abs().max());
436 }
437 return ret;
438 }
439
440 T min() const {
441 T ret = std::numeric_limits<T>::max();
442 for (int i = 0; i < size; i++) {
443 ret = std::min(ret, data[i]);
444 }
445 return ret;
446 }
447
448 T max() const {
449 T ret = std::numeric_limits<T>::min();
450 for (int i = 0; i < size; i++) {
451 ret = std::max(ret, data[i]);
452 }
453 return ret;
454 }
455
456 void print_abs_max_pos() const {
457 T ret = abs_max();
458 for (auto &ind : get_region()) {
459 if (abs(this->operator[](ind)) == ret) {
460 printf(" [%d, %d]\n", ind.i, ind.j);
461 }
462 }
463 }
464
465 template <typename TT = T,
466 typename std::enable_if_t<!std::is_class<TT>::value, int> = 0>
467 void print(std::string name = "", const char *temp = "%f") const {
468 if (name.size())
469 printf("%s[%dx%d]=", name.c_str(), res[0], res[1]);
470 printf("\n");
471 for (int j = res[1] - 1; j >= 0; j--) {
472 for (int i = 0; i < res[0]; i++) {
473 printf(temp, (*this)[i][j]);
474 printf(" ");
475 }
476 printf("\n");
477 }
478 printf("\n");
479 }
480
481 template <typename TT = T, typename TS = typename TT::ScalarType>
482 void print(std::string name = "", const char *temp = "%f") const {
483 if (name.size())
484 printf("%s[%dx%d]=", name.c_str(), res[0], res[1]);
485 printf("\n");
486 for (int j = res[1] - 1; j >= 0; j--) {
487 for (int i = 0; i < res[0]; i++) {
488 printf("(");
489 for (int k = 0; k < TT::D; k++) {
490 printf(temp, (*this)[i][j][k]);
491 if (k != TT::D - 1) {
492 printf(", ");
493 }
494 }
495 printf(") ");
496 }
497 printf("\n");
498 }
499 printf("\n");
500 }
501
502 size_t get_data_size() const {
503 return size * sizeof(T);
504 }
505
506 void set_pattern(int s) {
507 for (int i = 0; i < size; i++) {
508 data[i] = sinf(s * i + 231.0_f);
509 }
510 }
511
512 bool inside(int i, int j) const {
513 return 0 <= i && i < res[0] && 0 <= j && j < res[1];
514 }
515
516 bool inside(const Vector2i &pos) const {
517 return inside(pos[0], pos[1]);
518 }
519
520 bool inside(const Index2D &index) const {
521 return inside(index.i, index.j);
522 }
523
524 T sample(real x, real y) const {
525 x = clamp(x - storage_offset.x, 0.0_f, res[0] - 1.0_f - eps);
526 y = clamp(y - storage_offset.y, 0.0_f, res[1] - 1.0_f - eps);
527 int x_i = clamp(int(x), 0, res[0] - 2);
528 int y_i = clamp(int(y), 0, res[1] - 2);
529 real x_r = x - x_i;
530 real y_r = y - y_i;
531 return lerp(x_r, lerp(y_r, get(x_i, y_i), get(x_i, y_i + 1)),
532 lerp(y_r, get(x_i + 1, y_i), get(x_i + 1, y_i + 1)));
533 }
534
535 T sample(const Vector2 &v) const {
536 return sample(v.x, v.y);
537 }
538
539 T sample(const Index2D &v) const {
540 return sample(v.get_pos());
541 }
542
543 Vector2 get_storage_offset() const {
544 return storage_offset;
545 }
546
547 T sample_relative_coord(real x, real y) const {
548 x = x * res[0];
549 y = y * res[1];
550 return sample(x, y);
551 }
552
553 T sample_relative_coord(const Vector2 &vec) const {
554 real x = vec.x * res[0];
555 real y = vec.y * res[1];
556 return sample(x, y);
557 }
558
559 auto begin() const {
560 return data.cbegin();
561 }
562
563 auto end() const {
564 return data.cend();
565 }
566
567 auto begin() {
568 return data.begin();
569 }
570
571 auto end() {
572 return data.end();
573 }
574
575 T &operator[](const Vector2i &pos) {
576 return (*this)[pos.x][pos.y];
577 }
578
579 const T &operator[](const Vector2i &pos) const {
580 return (*this)[pos.x][pos.y];
581 }
582
583 T &operator[](const Index2D &index) {
584 return (*this)[index.i][index.j];
585 }
586
587 const T &operator[](const Index2D &index) const {
588 return (*this)[index.i][index.j];
589 }
590
591 Vector2i get_res() const {
592 return res;
593 }
594
595 int get_width() const {
596 return res[0];
597 }
598
599 int get_height() const {
600 return res[1];
601 }
602
603 bool empty() const {
604 return !(res[0] > 0 && res[1] > 0);
605 }
606
607 T get_average() const {
608 T sum(0);
609 for (int i = 0; i < res[0]; i++) {
610 for (int j = 0; j < res[1]; j++) {
611 sum += get(i, j);
612 }
613 }
614 return 1.0_f / size * sum;
615 }
616
617 bool inside(const Vector2 &pos, real tolerance = 1e-4f) const {
618 return (-tolerance <= pos.x && pos.x <= res[0] + tolerance &&
619 -tolerance <= pos.y && pos.y < res[1] + tolerance);
620 }
621
622 Region2D get_rasterization_region(Vector2 pos, int half_extent) const {
623 int x = (int)floor(pos.x - storage_offset.x);
624 int y = (int)floor(pos.y - storage_offset.y);
625 return Region2D(std::max(0, x - half_extent + 1),
626 std::min(res[0], x + half_extent + 1),
627 std::max(0, y - half_extent + 1),
628 std::min(res[1], y + half_extent + 1), storage_offset);
629 }
630
631 bool is_normal() const {
632 for (auto v : (*this)) {
633 if (!taichi::is_normal(v)) {
634 return false;
635 }
636 }
637 return true;
638 }
639
640 Array2D<T> rasterize(int width, int height) {
641 Array2D<T> out(Vector2i(width, height));
642 Vector2 actual_size;
643 if (storage_offset == Vector2(0.0_f, 0.0_f)) {
644 actual_size = Vector2(this->res[0] - 1, this->res[1] - 1);
645 } else {
646 actual_size = Vector2(this->res[0], this->res[1]);
647 }
648
649 Vector2 scale_factor = actual_size / res.cast<real>();
650
651 for (auto &ind : Region2D(0, res[0], 0, res[1], Vector2(0.5f, 0.5f))) {
652 Vector2 p = scale_factor * ind.get_pos();
653 out[ind] = sample(p);
654 }
655 return out;
656 }
657
658 Array2D<T> rasterize_scale(int width, int height, int scale) {
659 Array2D<T> out(Vector2i(width, height));
660 for (auto &ind : out.get_region()) {
661 out[ind] = (*this)[ind.i / scale][ind.j / scale];
662 }
663 return out;
664 }
665
666 const std::vector<T> &get_data() const {
667 return this->data;
668 }
669
670 static constexpr int get_dim() {
671 return 2;
672 }
673
674 void flip(int axis) {
675 if (axis == 0) {
676 for (int i = 0; i < res[0] / 2; i++) {
677 for (int j = 0; j < res[1]; j++) {
678 std::swap((*this)[i][j], (*this)[res[0] - 1 - i][j]);
679 }
680 }
681 } else {
682 for (int i = 0; i < res[0]; i++) {
683 for (int j = 0; j < res[1] / 2; j++) {
684 std::swap((*this)[i][j], (*this)[i][res[1] - 1 - j]);
685 }
686 }
687 }
688 }
689
690 // TODO: finally we are going to need a binary serializer
691
692 void write_to_disk(const std::string &fn) {
693 FILE *f = fopen(fn.c_str(), "wb");
694 fwrite(&res[0], sizeof(res[0]), 1, f);
695 fwrite(&res[1], sizeof(res[1]), 1, f);
696 fwrite(&storage_offset, sizeof(storage_offset), 1, f);
697 fwrite(&region, sizeof(region), 1, f);
698 fwrite(&data[0], sizeof(data[0]), size, f);
699 fclose(f);
700 }
701
702 bool read_from_disk(const std::string &fn) {
703 FILE *f = fopen(fn.c_str(), "rb");
704 if (f == nullptr) {
705 return false;
706 }
707 size_t ret;
708 ret = fread(&res[0], sizeof(res[0]), 1, f);
709 if (ret != 1) {
710 return false;
711 }
712 ret = fread(&res[1], sizeof(res[1]), 1, f);
713 if (ret != 1) {
714 return false;
715 }
716 ret = fread(&storage_offset, sizeof(storage_offset), 1, f);
717 if (ret != 1) {
718 return false;
719 }
720 ret = fread(&region, sizeof(region), 1, f);
721 if (ret != 1) {
722 return false;
723 }
724 initialize(res, T(0), storage_offset);
725 ret = (int)std::fread(&data[0], sizeof(data[0]), size, f);
726 if (ret != (std::size_t)size) {
727 return false;
728 }
729 fclose(f);
730 return true;
731 }
732
733 explicit ArrayND(const std::string &filename) {
734 load_image(filename);
735 }
736
737 void load_image(const std::string &filename, bool linearize = true);
738
739 void set_pixel(real x, real y, const T &pixel) {
740 x *= this->res[0];
741 y *= this->res[1];
742 x -= 0.5f;
743 y -= 0.5f;
744 int int_x = (int)x;
745 int int_y = (int)y;
746 if (int_x < 0 || int_x >= this->res[0] || int_y < 0 ||
747 int_y >= this->res[1])
748 return;
749 this->operator[](int_x)[int_y] = pixel;
750 }
751
752 T sample_as_texture(real x, real y, bool interp = true) {
753 x *= this->res[0];
754 y *= this->res[1];
755 x -= 0.5f;
756 y -= 0.5f;
757 x = clamp(x, 0.0_f, this->res[0] - 1.0_f);
758 y = clamp(y, 0.0_f, this->res[1] - 1.0_f);
759 int ix = clamp(int(x), 0, this->res[0] - 2);
760 int iy = clamp(int(y), 0, this->res[1] - 2);
761 if (!interp) {
762 x = real(ix);
763 y = real(iy);
764 }
765 T x_0 = lerp(y - iy, (*this)[ix][iy], (*this)[ix][iy + 1]);
766 T x_1 = lerp(y - iy, (*this)[ix + 1][iy], (*this)[ix + 1][iy + 1]);
767 return lerp(x - ix, x_0, x_1);
768 }
769
770 void write_as_image(const std::string &filename);
771
772 void write_text(const std::string &font_fn,
773 const std::string &content,
774 real size,
775 int dx,
776 int dy,
777 T color = T(1.0_f));
778};
779
780template <typename T>
781using Array2D = ArrayND<2, T>;
782
783template <typename T, typename P>
784inline Array2D<T> operator*(const P &b, const Array2D<T> &a) {
785 Array2D<T> o(a.res);
786 for (int i = 0; i < a.size; i++) {
787 o.data[i] = b * a.data[i];
788 }
789 return o;
790}
791
792template <typename T>
793inline void print(const Array2D<T> &arr) {
794 arr.print("");
795}
796
797} // namespace taichi
798